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#[cfg(test)]
304mod tests {
305    use super::*;
306    use crate::style::Style;
307
308    #[test]
309    fn test_segment_cell_length() {
310        let seg = Segment::new("Hello");
311        assert_eq!(seg.cell_length(), 5);
312    }
313
314    #[test]
315    fn test_segment_split() {
316        let seg = Segment::new("Hello World");
317        let (left, right) = seg.split(5);
318        assert_eq!(left.text, "Hello");
319        assert_eq!(right.unwrap().text, " World");
320    }
321
322    #[test]
323    fn test_segment_to_ansi() {
324        let style = Style::new().bold(true);
325        let seg = Segment::styled("Bold", style);
326        let ansi = seg.to_ansi();
327        assert!(ansi.contains("[1m"));
328        assert!(ansi.contains("Bold"));
329    }
330}