1use std::future::{Future, IntoFuture};
4use std::pin::Pin;
5
6use crate::{Locator, Pane, PaneCell, PaneSnapshot, Result};
7
8#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
10pub struct Rect {
11 pub row: u16,
13 pub col: u16,
15 pub rows: u16,
17 pub cols: u16,
19}
20
21impl Rect {
22 #[must_use]
24 pub const fn new(row: u16, col: u16, rows: u16, cols: u16) -> Self {
25 Self {
26 row,
27 col,
28 rows,
29 cols,
30 }
31 }
32}
33
34#[derive(Debug, Clone, PartialEq, Eq)]
36pub struct CapturedRegion {
37 pub rect: Rect,
39 pub text: String,
41 pub styled_cells: Option<Vec<PaneCell>>,
43 pub revision: u64,
45}
46
47#[derive(Debug, Clone)]
49#[must_use = "capture builders do nothing unless awaited"]
50pub struct CaptureBuilder {
51 source: CaptureSource,
52 preserve_style: bool,
53}
54
55#[derive(Debug, Clone)]
56enum CaptureSource {
57 Pane { pane: Pane, rect: Option<Rect> },
58 Locator(Locator),
59}
60
61impl CaptureBuilder {
62 pub(crate) fn pane(pane: Pane, rect: Option<Rect>) -> Self {
63 Self {
64 source: CaptureSource::Pane { pane, rect },
65 preserve_style: false,
66 }
67 }
68
69 pub(crate) fn locator(locator: Locator) -> Self {
70 Self {
71 source: CaptureSource::Locator(locator),
72 preserve_style: false,
73 }
74 }
75
76 pub const fn preserve_style(mut self, preserve: bool) -> Self {
78 self.preserve_style = preserve;
79 self
80 }
81
82 async fn run(self) -> Result<CapturedRegion> {
83 match self.source {
84 CaptureSource::Pane { pane, rect } => {
85 let snapshot = pane.snapshot().await?;
86 let rect = rect.unwrap_or_else(|| full_rect(&snapshot));
87 Ok(capture_from_snapshot(&snapshot, rect, self.preserve_style))
88 }
89 CaptureSource::Locator(locator) => {
90 let (snapshot, item) = locator.resolve_strict_with_wait().await?;
91 let rect = Rect::new(
92 item.text_match.start_row,
93 item.text_match.start_col,
94 item.text_match
95 .end_row
96 .saturating_sub(item.text_match.start_row)
97 .saturating_add(1),
98 item.text_match
99 .end_col
100 .saturating_sub(item.text_match.start_col),
101 );
102 Ok(capture_from_snapshot(&snapshot, rect, self.preserve_style))
103 }
104 }
105 }
106}
107
108impl IntoFuture for CaptureBuilder {
109 type Output = Result<CapturedRegion>;
110 type IntoFuture = Pin<Box<dyn Future<Output = Self::Output> + Send>>;
111
112 fn into_future(self) -> Self::IntoFuture {
113 Box::pin(self.run())
114 }
115}
116
117impl Pane {
118 pub fn capture_region(&self, rect: Rect) -> CaptureBuilder {
120 CaptureBuilder::pane(self.clone(), Some(rect))
121 }
122
123 pub fn screenshot(&self) -> CaptureBuilder {
125 CaptureBuilder::pane(self.clone(), None)
126 }
127}
128
129impl Locator {
130 pub async fn bounding_box(self) -> Result<Rect> {
132 let (_snapshot, item) = self.resolve_strict_with_wait().await?;
133 Ok(Rect::new(
134 item.text_match.start_row,
135 item.text_match.start_col,
136 item.text_match
137 .end_row
138 .saturating_sub(item.text_match.start_row)
139 .saturating_add(1),
140 item.text_match
141 .end_col
142 .saturating_sub(item.text_match.start_col),
143 ))
144 }
145
146 pub fn capture(self) -> CaptureBuilder {
148 CaptureBuilder::locator(self)
149 }
150
151 pub fn screenshot(self) -> CaptureBuilder {
153 self.capture()
154 }
155}
156
157fn capture_from_snapshot(
158 snapshot: &PaneSnapshot,
159 rect: Rect,
160 preserve_style: bool,
161) -> CapturedRegion {
162 let rect = clamp_rect(snapshot, rect);
163 let text = capture_text(snapshot, rect);
164 let styled_cells = preserve_style.then(|| capture_cells(snapshot, rect));
165 CapturedRegion {
166 rect,
167 text,
168 styled_cells,
169 revision: snapshot.revision,
170 }
171}
172
173fn full_rect(snapshot: &PaneSnapshot) -> Rect {
174 Rect::new(0, 0, snapshot.rows, snapshot.cols)
175}
176
177fn clamp_rect(snapshot: &PaneSnapshot, rect: Rect) -> Rect {
178 let row = rect.row.min(snapshot.rows);
179 let col = rect.col.min(snapshot.cols);
180 let rows = rect.rows.min(snapshot.rows.saturating_sub(row));
181 let cols = rect.cols.min(snapshot.cols.saturating_sub(col));
182 Rect::new(row, col, rows, cols)
183}
184
185fn capture_text(snapshot: &PaneSnapshot, rect: Rect) -> String {
186 (0..rect.rows)
187 .map(|offset| capture_row_text(snapshot, rect.row + offset, rect.col, rect.cols))
188 .collect::<Vec<_>>()
189 .join("\n")
190}
191
192fn capture_row_text(snapshot: &PaneSnapshot, row: u16, col: u16, cols: u16) -> String {
193 let mut text = String::new();
194 let end = col.saturating_add(cols).min(snapshot.cols);
195 for current_col in col..end {
196 let Some(cell) = snapshot.cell(row, current_col) else {
197 continue;
198 };
199 if !cell.is_padding() {
200 text.push_str(cell.text());
201 }
202 }
203 text.trim_end_matches(' ').to_owned()
204}
205
206fn capture_cells(snapshot: &PaneSnapshot, rect: Rect) -> Vec<PaneCell> {
207 let mut cells = Vec::new();
208 let end_row = rect.row.saturating_add(rect.rows).min(snapshot.rows);
209 let end_col = rect.col.saturating_add(rect.cols).min(snapshot.cols);
210 for row in rect.row..end_row {
211 for col in rect.col..end_col {
212 if let Some(cell) = snapshot.cell(row, col) {
213 cells.push(cell.clone());
214 }
215 }
216 }
217 cells
218}
219
220#[cfg(test)]
221mod tests {
222 use super::{capture_from_snapshot, Rect};
223 use crate::{PaneCell, PaneCursor, PaneGlyph, PaneSnapshot};
224
225 fn cell(text: &str) -> PaneCell {
226 PaneCell::new(PaneGlyph::new(text, 1))
227 }
228
229 #[test]
230 fn capture_region_clamps_out_of_bounds_rects() {
231 let snapshot = PaneSnapshot::new(
232 4,
233 2,
234 vec![
235 cell("a"),
236 cell("b"),
237 cell("c"),
238 cell("d"),
239 cell("e"),
240 cell("f"),
241 cell("g"),
242 cell("h"),
243 ],
244 PaneCursor::default(),
245 )
246 .expect("valid snapshot");
247
248 let capture = capture_from_snapshot(&snapshot, Rect::new(1, 2, 10, 10), false);
249
250 assert_eq!(capture.rect, Rect::new(1, 2, 1, 2));
251 assert_eq!(capture.text, "gh");
252 assert!(capture.styled_cells.is_none());
253 }
254
255 #[test]
256 fn styled_capture_preserves_row_major_cells_inside_clamped_region() {
257 let snapshot = PaneSnapshot::new(
258 3,
259 1,
260 vec![cell("x"), cell("y"), cell("z")],
261 PaneCursor::default(),
262 )
263 .expect("valid snapshot");
264
265 let capture = capture_from_snapshot(&snapshot, Rect::new(0, 1, 1, 9), true);
266
267 assert_eq!(capture.text, "yz");
268 let cells = capture.styled_cells.expect("styled cells");
269 assert_eq!(cells.len(), 2);
270 assert_eq!(cells[0].text(), "y");
271 assert_eq!(cells[1].text(), "z");
272 }
273}