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