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