rich/

live_render.rs

//! In-place live rendering.
//!
//! Port of `rich/live_render.py`. A [`LiveRender`] wraps a renderable, remembers
//! the shape (width × height) of its last render, and produces the terminal
//! control sequences to move the cursor back over that render — the mechanism a
//! `Live` display uses to redraw in place. The full `Live` loop (threading,
//! timing, stdout management) is deferred; this is its byte-parity-testable core.
//!
//! Content taller than the screen is cropped, ended with an ellipsis line
//! (the default) or left visible, per [`VerticalOverflow`].

use std::cell::Cell;

use crate::console::{Console, ConsoleOptions, Justify, Overflow};
use crate::control::{Control, ControlType};
use crate::protocol::Renderable;
use crate::segment::Segment;
use crate::style::Style;
use crate::text::Text;

/// What a live render does with content taller than the screen. Mirrors
/// upstream's `VerticalOverflowMethod`.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum VerticalOverflow {
    /// Keep only the lines that fit.
    Crop,
    /// Keep all but the last line that fits, then a centred `...` line in
    /// `live.ellipsis` (upstream's default).
    #[default]
    Ellipsis,
    /// Render every line.
    Visible,
}

/// Wraps a renderable for repeated in-place redraws. Mirrors
/// `rich.live_render.LiveRender`.
pub struct LiveRender {
    renderable: Box<dyn Renderable>,
    style: Option<Style>,
    vertical_overflow: VerticalOverflow,
    /// `(width, height)` of the last render, or `None` before the first.
    shape: Cell<Option<(usize, usize)>>,
}

impl LiveRender {
    pub fn new(renderable: Box<dyn Renderable>) -> Self {
        LiveRender {
            renderable,
            style: None,
            vertical_overflow: VerticalOverflow::Ellipsis,
            shape: Cell::new(None),
        }
    }

    /// Apply a style across the whole live render.
    pub fn style(mut self, style: Style) -> Self {
        self.style = Some(style);
        self
    }

    /// What to do with content taller than the screen (upstream
    /// `vertical_overflow`, default [`VerticalOverflow::Ellipsis`]).
    pub fn vertical_overflow(mut self, vertical_overflow: VerticalOverflow) -> Self {
        self.vertical_overflow = vertical_overflow;
        self
    }

    /// Set [`vertical_overflow`](Self::vertical_overflow) in place.
    pub fn set_vertical_overflow(&mut self, vertical_overflow: VerticalOverflow) {
        self.vertical_overflow = vertical_overflow;
    }

    /// Replace the wrapped renderable (the shape carries over until the next
    /// render). Port of `LiveRender.set_renderable`.
    pub fn set_renderable(&mut self, renderable: Box<dyn Renderable>) {
        self.renderable = renderable;
    }

    /// The wrapped renderable. Upstream's `LiveRender.renderable`.
    pub fn renderable(&self) -> &dyn Renderable {
        self.renderable.as_ref()
    }

    /// The height of the last render, `0` before any. Port of
    /// `LiveRender.last_render_height`.
    pub fn last_render_height(&self) -> usize {
        self.shape.get().map_or(0, |(_, height)| height)
    }

    /// Control codes to move the cursor to the start of the previous render,
    /// erasing each line. Port of `LiveRender.position_cursor`.
    pub fn position_cursor(&self) -> Control {
        position_cursor(self.shape.get())
    }

    /// Control codes to clear the render and restore the cursor to where it was
    /// before it. Port of `LiveRender.restore_cursor`.
    pub fn restore_cursor(&self) -> Control {
        restore_cursor(self.shape.get())
    }
}

/// [`LiveRender::position_cursor`] for a render of `shape`.
pub(crate) fn position_cursor(shape: Option<(usize, usize)>) -> Control {
    match shape {
        Some((_, height)) => {
            let mut codes = vec![ControlType::CarriageReturn, ControlType::EraseInLine(2)];
            for _ in 0..height.saturating_sub(1) {
                codes.push(ControlType::CursorUp(1));
                codes.push(ControlType::EraseInLine(2));
            }
            Control::new(&codes)
        }
        None => Control::new(&[]),
    }
}

/// [`LiveRender::restore_cursor`] for a render of `shape`.
pub(crate) fn restore_cursor(shape: Option<(usize, usize)>) -> Control {
    match shape {
        Some((_, height)) => {
            let mut codes = vec![ControlType::CarriageReturn];
            for _ in 0..height {
                codes.push(ControlType::CursorUp(1));
                codes.push(ControlType::EraseInLine(2));
            }
            Control::new(&codes)
        }
        None => Control::new(&[]),
    }
}

impl Renderable for LiveRender {
    fn rich_render(&self, console: &Console, options: &ConsoleOptions) -> Vec<Segment> {
        let (segments, shape) = render(
            self.renderable.as_ref(),
            self.style.as_ref(),
            self.vertical_overflow,
            console,
            options,
        );
        // Remember the shape for the next `position_cursor`/`restore_cursor`.
        self.shape.set(Some(shape));
        segments
    }

    /// Upstream yields no line break after the last line, so a printed
    /// `LiveRender` leaves the cursor at its end (where `position_cursor`
    /// expects it).
    fn print_end_line(&self, _options: &ConsoleOptions) -> bool {
        false
    }
}

/// `LiveRender.__rich_console__` over a borrowed renderable: the segments,
/// and the `(width, height)` shape the render leaves.
pub(crate) fn render(
    renderable: &dyn Renderable,
    style: Option<&Style>,
    vertical_overflow: VerticalOverflow,
    console: &Console,
    options: &ConsoleOptions,
) -> (Vec<Segment>, (usize, usize)) {
    // Render unpadded lines (`pad=False`), applying the live style if set.
    let mut lines = console.render_lines(renderable, options, false);
    if let Some(style) = style {
        for line in &mut lines {
            *line = Segment::apply_style(line, style);
        }
    }

    // Content taller than the screen (`options.size.height`).
    let screen_height = options.size.height;
    if lines.len() > screen_height {
        match vertical_overflow {
            VerticalOverflow::Crop => lines.truncate(screen_height),
            VerticalOverflow::Ellipsis => {
                // `lines[: height - 1]`: at height 0 that is Python's
                // `lines[:-1]`, every line but the last.
                let keep = screen_height
                    .checked_sub(1)
                    .unwrap_or_else(|| lines.len() - 1);
                lines.truncate(keep);
                let overflow_text = Text::styled("...", "live.ellipsis")
                    .overflow(Overflow::Crop)
                    .justify(Justify::Center);
                lines.push(console.render(&overflow_text, None));
            }
            VerticalOverflow::Visible => {}
        }
    }

    let height = lines.len();
    let width = lines
        .iter()
        .map(|line| line.iter().map(Segment::cell_length).sum::<usize>())
        .max()
        .unwrap_or(0);

    let mut segments = Vec::new();
    let last = height.saturating_sub(1);
    for (index, line) in lines.into_iter().enumerate() {
        segments.extend(line);
        if index != last {
            segments.push(Segment::line());
        }
    }
    (segments, (width, height))
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::color::ColorSystem;
    use crate::text::Text;

    fn console() -> Console {
        Console::builder()
            .force_terminal(true)
            .color_system(Some(ColorSystem::Truecolor))
            .width(20)
            .no_color(false)
            .build()
    }

    #[test]
    fn renders_content_and_control_codes() {
        let live = LiveRender::new(Box::new(Text::new("line one\nline two\nline three")));
        let console = console();
        // Rendering sets the shape.
        assert_eq!(
            console.render_to_string(&live),
            "line one\nline two\nline three"
        );
        // Captured from real rich 15.0.0 (height 3).
        assert_eq!(
            live.position_cursor().as_str(),
            "\r\x1b[2K\x1b[1A\x1b[2K\x1b[1A\x1b[2K"
        );
        assert_eq!(
            live.restore_cursor().as_str(),
            "\r\x1b[1A\x1b[2K\x1b[1A\x1b[2K\x1b[1A\x1b[2K"
        );
    }

    #[test]
    fn single_line_shape() {
        let live = LiveRender::new(Box::new(Text::new("solo")));
        let console = console();
        assert_eq!(console.render_to_string(&live), "solo");
        assert_eq!(live.position_cursor().as_str(), "\r\x1b[2K");
        assert_eq!(live.restore_cursor().as_str(), "\r\x1b[1A\x1b[2K");
    }

    #[test]
    fn no_control_before_first_render() {
        let live = LiveRender::new(Box::new(Text::new("x")));
        assert_eq!(live.position_cursor().as_str(), "");
        assert_eq!(live.restore_cursor().as_str(), "");
    }
}