Skip to main content

rusty_rich/
segment.rs

1//! Segment — styled text unit. Equivalent to Rich's `segment.py`.
2//!
3//! A `Segment` is the smallest unit of output: a piece of text with an
4//! associated style and optional control codes.
5
6use std::fmt;
7
8use crate::style::Style;
9
10// ---------------------------------------------------------------------------
11// ControlType
12// ---------------------------------------------------------------------------
13
14/// Non-printable control codes (equivalent to Rich's `ControlType`).
15#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
16pub enum ControlType {
17    Bell,
18    CarriageReturn,
19    Home,
20    Clear,
21    ShowCursor,
22    HideCursor,
23    EnableAltScreen,
24    DisableAltScreen,
25    CursorUp,
26    CursorDown,
27    CursorForward,
28    CursorBackward,
29    CursorMoveToColumn,
30    CursorMoveTo,
31    EraseInLine,
32    SetWindowTitle,
33}
34
35impl ControlType {
36    /// Get the ANSI escape sequence for this control type.
37    pub fn to_ansi(&self, params: &[i32]) -> String {
38        match self {
39            Self::Bell => "\x07".into(),
40            Self::CarriageReturn => "\r".into(),
41            Self::Home => "\x1b[H".into(),
42            Self::Clear => "\x1b[2J".into(),
43            Self::ShowCursor => "\x1b[?25h".into(),
44            Self::HideCursor => "\x1b[?25l".into(),
45            Self::EnableAltScreen => "\x1b[?1049h".into(),
46            Self::DisableAltScreen => "\x1b[?1049l".into(),
47            Self::CursorUp => {
48                let n = params.first().copied().unwrap_or(1);
49                format!("\x1b[{n}A")
50            }
51            Self::CursorDown => {
52                let n = params.first().copied().unwrap_or(1);
53                format!("\x1b[{n}B")
54            }
55            Self::CursorForward => {
56                let n = params.first().copied().unwrap_or(1);
57                format!("\x1b[{n}C")
58            }
59            Self::CursorBackward => {
60                let n = params.first().copied().unwrap_or(1);
61                format!("\x1b[{n}D")
62            }
63            Self::CursorMoveToColumn => {
64                let col = params.first().copied().unwrap_or(0);
65                format!("\x1b[{col}G")
66            }
67            Self::CursorMoveTo => {
68                let row = params.first().copied().unwrap_or(0);
69                let col = params.get(1).copied().unwrap_or(0);
70                format!("\x1b[{row};{col}H")
71            }
72            Self::EraseInLine => {
73                let mode = params.first().copied().unwrap_or(0);
74                format!("\x1b[{mode}K")
75            }
76            Self::SetWindowTitle => {
77                let title: String = params
78                    .iter()
79                    .map(|n| char::from(*n as u8))
80                    .collect();
81                format!("\x1b]0;{title}\x07")
82            }
83        }
84    }
85}
86
87// ---------------------------------------------------------------------------
88// ControlCode
89// ---------------------------------------------------------------------------
90
91#[derive(Debug, Clone, PartialEq, Eq, Hash)]
92pub enum ControlCode {
93    Simple(ControlType),
94    WithInt(ControlType, i32),
95    WithTwoInts(ControlType, i32, i32),
96    WithString(ControlType, String),
97}
98
99// ---------------------------------------------------------------------------
100// Segment
101// ---------------------------------------------------------------------------
102
103/// A piece of text with an associated style.
104///
105/// Segments are produced during rendering and ultimately converted to strings
106/// for terminal output.
107#[derive(Debug, Clone, PartialEq)]
108pub struct Segment {
109    pub text: String,
110    pub style: Option<Style>,
111    pub control: Option<ControlCode>,
112}
113
114impl Segment {
115    /// Create a new segment with just text.
116    pub fn new(text: impl Into<String>) -> Self {
117        Self {
118            text: text.into(),
119            style: None,
120            control: None,
121        }
122    }
123
124    /// Create a segment with text and style.
125    pub fn styled(text: impl Into<String>, style: Style) -> Self {
126        Self {
127            text: text.into(),
128            style: Some(style),
129            control: None,
130        }
131    }
132
133    /// Create a control-only segment.
134    pub fn control(code: ControlCode) -> Self {
135        Self {
136            text: String::new(),
137            style: None,
138            control: Some(code),
139        }
140    }
141
142    /// Create a newline segment.
143    pub fn line() -> Self {
144        Self::new("\n")
145    }
146
147    /// Get the cell length of this segment's text (using Unicode width).
148    pub fn cell_length(&self) -> usize {
149        if self.control.is_some() {
150            return 0;
151        }
152        unicode_width::UnicodeWidthStr::width(self.text.as_str())
153    }
154
155    /// Check if this segment has any text content.
156    pub fn is_empty(&self) -> bool {
157        self.text.is_empty() && self.control.is_none()
158    }
159
160    /// Split this segment at the given offset, returning two segments.
161    /// The first goes up to (but not including) `offset`, the second from
162    /// `offset` to end.
163    pub fn split(&self, offset: usize) -> (Segment, Option<Segment>) {
164        if offset == 0 {
165            return (Segment::new(""), Some(self.clone()));
166        }
167        let cell_len = self.cell_length();
168        if offset >= cell_len {
169            return (self.clone(), None);
170        }
171
172        // Find the byte position at which the cell count reaches `offset`
173        let mut cell_count = 0usize;
174        let mut byte_pos = 0usize;
175        for (i, ch) in self.text.char_indices() {
176            let w = unicode_width::UnicodeWidthChar::width(ch).unwrap_or(0);
177            if cell_count + w > offset {
178                break;
179            }
180            cell_count += w;
181            byte_pos = i + ch.len_utf8();
182        }
183
184        let left = Segment {
185            text: self.text[..byte_pos].to_string(),
186            style: self.style.clone(),
187            control: self.control.clone(),
188        };
189        let right = Segment {
190            text: self.text[byte_pos..].to_string(),
191            style: self.style.clone(),
192            control: self.control.clone(),
193        };
194        (left, Some(right))
195    }
196
197    /// Return the ANSI string for this segment (style + text + reset).
198    pub fn to_ansi(&self) -> String {
199        if let Some(ref code) = self.control {
200            return match code {
201                ControlCode::Simple(ct) => ct.to_ansi(&[]),
202                ControlCode::WithInt(ct, a) => ct.to_ansi(&[*a]),
203                ControlCode::WithTwoInts(ct, a, b) => ct.to_ansi(&[*a, *b]),
204                ControlCode::WithString(ct, s) => {
205                    let params: Vec<i32> = s.bytes().map(|b| b as i32).collect();
206                    ct.to_ansi(&params)
207                }
208            };
209        }
210
211        let style_ansi = self.style.as_ref().map(|s| s.to_ansi()).unwrap_or_default();
212        let reset = self.style.as_ref().map(|s| s.reset_ansi()).unwrap_or("");
213
214        if style_ansi.is_empty() {
215            self.text.clone()
216        } else {
217            format!("{style_ansi}{}{reset}", self.text)
218        }
219    }
220}
221
222impl fmt::Display for Segment {
223    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
224        write!(f, "{}", self.to_ansi())
225    }
226}
227
228// ---------------------------------------------------------------------------
229// Segments — a collection of segments
230// ---------------------------------------------------------------------------
231
232/// A collection of `Segment`s, with convenience methods.
233#[derive(Debug, Clone, Default)]
234pub struct Segments {
235    pub segments: Vec<Segment>,
236}
237
238impl Segments {
239    pub fn new() -> Self {
240        Self {
241            segments: Vec::new(),
242        }
243    }
244
245    pub fn push(&mut self, seg: Segment) {
246        self.segments.push(seg);
247    }
248
249    pub fn extend(&mut self, other: impl IntoIterator<Item = Segment>) {
250        self.segments.extend(other);
251    }
252
253    /// Render all segments to an ANSI string.
254    pub fn to_ansi(&self) -> String {
255        let mut out = String::new();
256        for seg in &self.segments {
257            out.push_str(&seg.to_ansi());
258        }
259        out
260    }
261
262    /// Get the total cell width.
263    pub fn cell_len(&self) -> usize {
264        self.segments.iter().map(Segment::cell_length).sum()
265    }
266}
267
268impl fmt::Display for Segments {
269    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
270        write!(f, "{}", self.to_ansi())
271    }
272}
273
274impl From<Vec<Segment>> for Segments {
275    fn from(segments: Vec<Segment>) -> Self {
276        Self { segments }
277    }
278}
279
280impl IntoIterator for Segments {
281    type Item = Segment;
282    type IntoIter = std::vec::IntoIter<Segment>;
283
284    fn into_iter(self) -> Self::IntoIter {
285        self.segments.into_iter()
286    }
287}
288
289// ---------------------------------------------------------------------------
290// Utility: line()
291// ---------------------------------------------------------------------------
292
293/// Helper: create a newline segment.
294pub fn line() -> Segment {
295    Segment::line()
296}
297
298/// Helper: create a space segment.
299pub fn space(count: usize) -> Segment {
300    Segment::new(" ".repeat(count))
301}
302
303// ---------------------------------------------------------------------------
304// Segment collection utilities
305// ---------------------------------------------------------------------------
306
307impl Segments {
308    /// Combine adjacent segments that have the same style.
309    pub fn simplify(&self) -> Segments {
310        let mut result: Vec<Segment> = Vec::new();
311        for seg in &self.segments {
312            if let Some(last) = result.last_mut() {
313                if last.style == seg.style && last.control.is_none() && seg.control.is_none() {
314                    last.text.push_str(&seg.text);
315                    continue;
316                }
317            }
318            result.push(seg.clone());
319        }
320        Segments { segments: result }
321    }
322}
323
324/// Split an iterable of segments into lines at newline boundaries.
325pub fn split_lines(segments: &[Segment]) -> Vec<Vec<Segment>> {
326    let mut lines: Vec<Vec<Segment>> = Vec::new();
327    let mut current: Vec<Segment> = Vec::new();
328    for seg in segments {
329        if seg.text == "\n" && seg.style.is_none() && seg.control.is_none() {
330            lines.push(std::mem::take(&mut current));
331        } else if seg.text.contains('\n') && seg.style.is_none() && seg.control.is_none() {
332            let parts: Vec<&str> = seg.text.split('\n').collect();
333            for (i, part) in parts.iter().enumerate() {
334                if i > 0 {
335                    lines.push(std::mem::take(&mut current));
336                }
337                if !part.is_empty() {
338                    current.push(Segment::new(*part));
339                }
340            }
341        } else {
342            current.push(seg.clone());
343        }
344    }
345    if !current.is_empty() {
346        lines.push(current);
347    }
348    lines
349}
350
351/// Remove all styles from segments, returning plain text only.
352pub fn strip_styles(segments: &[Segment]) -> String {
353    let mut out = String::new();
354    for seg in segments {
355        if seg.control.is_none() {
356            out.push_str(&seg.text);
357        }
358    }
359    out
360}
361
362/// Remove link IDs and URLs from all segment styles.
363pub fn strip_links(segments: &[Segment]) -> Vec<Segment> {
364    segments
365        .iter()
366        .map(|seg| {
367            let mut s = seg.clone();
368            if let Some(ref style) = seg.style {
369                let mut new_style = style.clone();
370                new_style.link_id = 0;
371                new_style.link = None;
372                s.style = Some(new_style);
373            }
374            s
375        })
376        .collect()
377}
378
379/// Align lines to the top of a region of given height.
380pub fn align_top(
381    lines: &[Vec<Segment>],
382    _width: usize,
383    height: usize,
384    _style: Option<&Style>,
385) -> Vec<Vec<Segment>> {
386    let blank_line = vec![Segment::new(" ".repeat(_width))];
387    let mut result: Vec<Vec<Segment>> = lines.to_vec();
388    while result.len() < height {
389        result.push(blank_line.clone());
390    }
391    result.truncate(height);
392    result
393}
394
395/// Align lines to the middle of a region of given height.
396pub fn align_middle(
397    lines: &[Vec<Segment>],
398    _width: usize,
399    height: usize,
400    _style: Option<&Style>,
401) -> Vec<Vec<Segment>> {
402    let blank_line = vec![Segment::new(" ".repeat(_width))];
403    let top_pad = (height.saturating_sub(lines.len())) / 2;
404    let mut result: Vec<Vec<Segment>> = Vec::new();
405    for _ in 0..top_pad {
406        result.push(blank_line.clone());
407    }
408    result.extend(lines.iter().cloned());
409    while result.len() < height {
410        result.push(blank_line.clone());
411    }
412    result.truncate(height);
413    result
414}
415
416/// Align lines to the bottom of a region of given height.
417pub fn align_bottom(
418    lines: &[Vec<Segment>],
419    _width: usize,
420    height: usize,
421    _style: Option<&Style>,
422) -> Vec<Vec<Segment>> {
423    let blank_line = vec![Segment::new(" ".repeat(_width))];
424    let bottom_pad = height.saturating_sub(lines.len());
425    let mut result: Vec<Vec<Segment>> = Vec::new();
426    for _ in 0..bottom_pad {
427        result.push(blank_line.clone());
428    }
429    result.extend(lines.iter().cloned());
430    result.truncate(height);
431    result
432}
433
434/// Divide segments at the given cell offsets.
435pub fn divide(segments: &[Segment], cuts: &[usize]) -> Vec<Vec<Segment>> {
436    let mut result: Vec<Vec<Segment>> = Vec::new();
437    let mut remaining = segments.to_vec();
438    let mut offset = 0usize;
439
440    for &cut in cuts {
441        let mut chunk: Vec<Segment> = Vec::new();
442        let target = cut.saturating_sub(offset);
443
444        let mut chunk_cells = 0usize;
445        while chunk_cells < target && !remaining.is_empty() {
446            let seg = remaining.remove(0);
447            let seg_len = seg.cell_length();
448            if chunk_cells + seg_len <= target {
449                chunk_cells += seg_len;
450                chunk.push(seg);
451            } else {
452                let split_at = target - chunk_cells;
453                let (left, right) = seg.split(split_at);
454                chunk.push(left);
455                if let Some(r) = right {
456                    remaining.insert(0, r);
457                }
458                chunk_cells = target;
459            }
460        }
461        result.push(chunk);
462        offset = cut;
463    }
464
465    if !remaining.is_empty() {
466        result.push(remaining);
467    }
468
469    result
470}
471
472/// Set segments to an exact width and height, padding/truncating as needed.
473pub fn set_shape(
474    lines: &[Vec<Segment>],
475    width: usize,
476    height: usize,
477    _style: Option<&Style>,
478) -> Vec<Vec<Segment>> {
479    let blank_line = vec![Segment::new(" ".repeat(width))];
480    let mut result: Vec<Vec<Segment>> = Vec::new();
481
482    for line in lines.iter().take(height) {
483        let cell_len: usize = line.iter().map(|s| s.cell_length()).sum();
484        let mut new_line = line.clone();
485        if cell_len < width {
486            new_line.push(Segment::new(" ".repeat(width - cell_len)));
487        } else if cell_len > width {
488            let mut truncated = Vec::new();
489            let mut count = 0usize;
490            for seg in line {
491                let seg_len = seg.cell_length();
492                if count + seg_len <= width {
493                    truncated.push(seg.clone());
494                    count += seg_len;
495                } else if count < width {
496                    let (left, _) = seg.split(width - count);
497                    truncated.push(left);
498                    break;
499                }
500            }
501            new_line = truncated;
502        }
503        result.push(new_line);
504    }
505
506    while result.len() < height {
507        result.push(blank_line.clone());
508    }
509
510    result
511}
512
513/// Filter segments, keeping only control codes if `is_control` is true,
514/// or only non-control segments if `is_control` is false.
515pub fn filter_control(segments: &[Segment], is_control: bool) -> Vec<Segment> {
516    segments
517        .iter()
518        .filter(|seg| seg.control.is_some() == is_control)
519        .cloned()
520        .collect()
521}
522
523/// Get the total cell length of a line of segments.
524pub fn get_line_length(line: &[Segment]) -> usize {
525    line.iter().map(|s| s.cell_length()).sum()
526}
527
528#[cfg(test)]
529mod tests {
530    use super::*;
531    use crate::style::Style;
532
533    #[test]
534    fn test_segment_cell_length() {
535        let seg = Segment::new("Hello");
536        assert_eq!(seg.cell_length(), 5);
537    }
538
539    #[test]
540    fn test_segment_split() {
541        let seg = Segment::new("Hello World");
542        let (left, right) = seg.split(5);
543        assert_eq!(left.text, "Hello");
544        assert_eq!(right.unwrap().text, " World");
545    }
546
547    #[test]
548    fn test_segment_to_ansi() {
549        let style = Style::new().bold(true);
550        let seg = Segment::styled("Bold", style);
551        let ansi = seg.to_ansi();
552        assert!(ansi.contains("[1m"));
553        assert!(ansi.contains("Bold"));
554    }
555}