rs-rich-cli 0.0.26

//! `rich chart`: a chart from CSV, JSON or stdin (0.0.15 workstream 5). Not
//! upstream: a binary-boundary convenience that composes `rich_ext::chart`;
//! see `docs/PORTING.md`.
//!
//! Reading the data into columns lives here rather than in `rich-ext`: CSV
//! goes through `rs-rich-data`'s port of Python's `csv` sniffer and reader
//! (the one `--csv` uses), and picking the columns from `--x` and `--y` is
//! command routing. The charts themselves are `rich_ext::chart`'s, unchanged.
use super::*;

use rich::table::Cell as TableCell;
use rich_ext::chart::{Bar, BarChart, BarGroup, Heatmap, LineChart, PieChart, Series, Sparkline};

/// `--kind`.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub(crate) enum ChartKind {
    /// One sparkline per series.
    Spark,
    /// A bar per row, labelled by `--x`.
    Bar,
    /// Each series as a line against `--x` (or the row number). The default.
    #[default]
    Line,
    /// Each series as points against `--x` (or the row number).
    Scatter,
    /// Rows by series, each value a shade.
    Heatmap,
    /// One series as slices of a disc, labelled by `--x`.
    Pie,
    /// A bar per row, labelled by `--x`, each series a part of it.
    Stacked,
}

impl ChartKind {
    fn parse(value: &str) -> Result<Self, String> {
        match value {
            "spark" | "sparkline" => Ok(Self::Spark),
            "bar" => Ok(Self::Bar),
            "line" => Ok(Self::Line),
            "scatter" => Ok(Self::Scatter),
            "heatmap" => Ok(Self::Heatmap),
            "pie" => Ok(Self::Pie),
            "stacked" => Ok(Self::Stacked),
            other => Err(format!(
                "unknown chart kind {other:?} (spark, bar, line, scatter, heatmap, pie, stacked)"
            )),
        }
    }

    /// Whether `--x` is a number (a position) rather than a label.
    fn numeric_x(self) -> bool {
        matches!(self, Self::Line | Self::Scatter)
    }
}

/// The options of `rich chart`.
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub(crate) struct ChartOptions {
    /// `--kind`; `None` when not given (a line chart).
    kind: Option<ChartKind>,
    /// `--x COLUMN`.
    x: Option<String>,
    /// `--y COLUMN`, repeatable: one series each.
    y: Vec<String>,
}

impl ChartOptions {
    /// Consume one of these options; anything else stays with the main parser.
    pub(crate) fn parse_option<'a>(
        &mut self,
        arg: &str,
        rest: &mut impl Iterator<Item = &'a String>,
    ) -> Result<bool, String> {
        match arg {
            "--kind" => {
                let value = rest
                    .next()
                    .ok_or("--kind requires spark, bar, line, scatter, heatmap, pie or stacked")?;
                self.kind = Some(ChartKind::parse(value)?);
            }
            "--x" => {
                let value = rest.next().ok_or("--x requires a COLUMN")?;
                if self.x.is_some() {
                    return Err("--x can be given once; a chart has one x column".into());
                }
                self.x = Some(value.clone());
            }
            "--y" => {
                let value = rest.next().ok_or("--y requires a COLUMN")?;
                self.y.push(value.clone());
            }
            _ => return Ok(false),
        }
        Ok(true)
    }

    /// Each option given with the commands it applies to, for the "only has
    /// an effect" check.
    pub(crate) fn given(&self) -> Vec<(&'static str, &'static [&'static str])> {
        [
            ("--kind", self.kind.is_some()),
            ("--x", self.x.is_some()),
            ("--y", !self.y.is_empty()),
        ]
        .into_iter()
        .filter(|(_, given)| *given)
        .map(|(flag, _)| (flag, &["chart"][..]))
        .collect()
    }
}

type Failure = (ExitClass, String);

/// One value read from the data, before a chart asks for it as a number.
#[derive(Debug, Clone, PartialEq)]
enum Value {
    /// An empty CSV cell, a JSON `null` or a key a record lacks: a gap.
    Missing,
    Number(f64),
    Text(String),
}

impl Value {
    /// The value as a label: a number as written, a gap as nothing.
    fn label(&self) -> String {
        match self {
            Value::Missing => String::new(),
            Value::Number(n) => n.to_string(),
            Value::Text(text) => text.clone(),
        }
    }

    /// The value as a number: `Ok(NaN)` for a gap, `Err` with why for text
    /// that does not read as one, or reads as one no chart can place
    /// (`1e400`, `nan`, `inf`).
    fn number(&self) -> Result<f64, String> {
        let (n, text) = match self {
            Value::Missing => return Ok(f64::NAN),
            Value::Number(n) => (*n, None),
            Value::Text(text) => {
                let trimmed = text.trim();
                if trimmed.is_empty() {
                    return Ok(f64::NAN);
                }
                match trimmed.parse::<f64>() {
                    Ok(n) => (n, Some(text.as_str())),
                    Err(_) => return Err(format!("{text:?} is not a number")),
                }
            }
        };
        if n.is_finite() {
            Ok(n)
        } else {
            let text = text.map_or_else(|| n.to_string(), str::to_string);
            Err(format!("{text:?} is not a finite number"))
        }
    }
}

/// The data as named columns of rows.
#[derive(Debug, Clone, PartialEq)]
struct Data {
    columns: Vec<String>,
    rows: Vec<Vec<Value>>,
}

impl Data {
    fn cell(&self, row: usize, column: usize) -> &Value {
        self.rows[row].get(column).unwrap_or(&Value::Missing)
    }

    /// The column `name` names: its header exactly, else a 1-based position.
    fn column(&self, name: &str) -> Result<usize, Failure> {
        if let Some(index) = self.columns.iter().position(|c| c == name) {
            return Ok(index);
        }
        if let Ok(n) = name.parse::<usize>()
            && (1..=self.columns.len()).contains(&n)
        {
            return Ok(n - 1);
        }
        let shown: Vec<String> = self.columns.iter().map(|c| format!("{c:?}")).collect();
        Err((
            ExitClass::Data,
            format!(
                "no column {name:?}; the columns are {}",
                if shown.is_empty() {
                    "none".to_string()
                } else {
                    shown.join(", ")
                }
            ),
        ))
    }

    /// Whether every value in `column` reads as a number (gaps aside), and
    /// at least one is there. `inf` and `1e400` count, so a column of
    /// numbers with one of them is charted and refused at that row.
    fn is_numeric(&self, column: usize) -> bool {
        let mut any = false;
        for row in 0..self.rows.len() {
            match self.cell(row, column) {
                Value::Missing => {}
                Value::Number(_) => any = true,
                Value::Text(text) if text.trim().is_empty() => {}
                Value::Text(text) => match text.trim().parse::<f64>() {
                    Ok(_) => any = true,
                    Err(_) => return false,
                },
            }
        }
        any
    }

    /// `column` as numbers; text that is not one is an error naming the row
    /// (counted from 1, after any header) and the column.
    fn numbers(&self, column: usize) -> Result<Vec<f64>, Failure> {
        (0..self.rows.len())
            .map(|row| {
                self.cell(row, column).number().map_err(|why| {
                    (
                        ExitClass::Data,
                        format!("row {}, column {:?}: {why}", row + 1, self.columns[column]),
                    )
                })
            })
            .collect()
    }
}

/// `rich chart [FILE]`: read the data, pick the columns, draw the chart.
pub(crate) fn chart(cli: &Cli) -> Result<Box<dyn Renderable>, Failure> {
    let resource = cli.resource.as_deref().unwrap_or("-");
    let content = if is_url(resource) {
        fetch_url(resource, cli.extensions.encoding)
            .map(|(content, _)| content)
            .map_err(|err| (ExitClass::Input, err))?
    } else {
        read_resource(Some(resource), cli.extensions.encoding).map_err(|err| {
            let name = if resource == "-" { "<stdin>" } else { resource };
            (ExitClass::Input, format!("cannot read {name}: {err}"))
        })?
    };
    // `--sanitize` (the default for a URL): column names and labels show
    // terminal controls instead of sending them.
    let content = if cli.sanitize {
        rich_ext::sanitize_terminal_controls(&content)
    } else {
        content
    };
    let name = if resource == "-" { "<stdin>" } else { resource };
    let data =
        parse(&content, resource).map_err(|err| (ExitClass::Data, format!("{name}: {err}")))?;
    build(&data, &cli.chart)
}

/// Read `content` as JSON or JSON Lines, CSV/TSV, or bare numbers: by the
/// resource's extension, else by its first character.
fn parse(content: &str, resource: &str) -> Result<Data, String> {
    let extension = resource
        .rsplit(['/', '\\'])
        .next()
        .and_then(|name| name.rsplit_once('.'))
        .map(|(_, ext)| ext.to_ascii_lowercase());
    // A URL's text keeps its byte-order mark; files and stdin lose it on
    // reading. Every parser gets the text without it.
    let content = content.trim_start_matches('\u{feff}');
    let trimmed = content.trim_start();
    let data = match extension.as_deref() {
        Some("json") => parse_json(content)?,
        // One record per line, even when there is only one.
        Some("jsonl" | "ndjson") => json_to_data(json_lines(content, None)?)?,
        Some("csv") => parse_csv(content, Some(','))?,
        Some("tsv") => parse_csv(content, Some('\t'))?,
        _ if trimmed.starts_with(['[', '{']) => parse_json(content)?,
        _ => match parse_numbers(content) {
            Some(data) => data,
            None => parse_csv(content, None)?,
        },
    };
    if data.rows.is_empty() {
        return Err("no rows to chart".into());
    }
    Ok(data)
}

/// Whitespace-separated numbers (`seq 10 | rich chart`): one column, `value`.
fn parse_numbers(content: &str) -> Option<Data> {
    let rows: Option<Vec<Vec<Value>>> = content
        .split_whitespace()
        .map(|token| {
            token.parse::<f64>().ok().map(|n| {
                // Kept as written, so a refusal quotes it.
                vec![if n.is_finite() {
                    Value::Number(n)
                } else {
                    Value::Text(token.to_string())
                }]
            })
        })
        .collect();
    let rows = rows?;
    (!rows.is_empty()).then(|| Data {
        columns: vec!["value".into()],
        rows,
    })
}

/// CSV or TSV, the dialect sniffed as `--csv` sniffs it (falling back to
/// `fallback`, or a comma). The first row is a header unless every cell in it
/// is a number; without one the columns are named `1`, `2`, …
fn parse_csv(content: &str, fallback: Option<char>) -> Result<Data, String> {
    let sample: String = content.chars().take(1024).collect();
    let dialect = sniff(&sample, Some(&[',', '\t', '|', ';']))
        .unwrap_or_else(|| Dialect::excel(fallback.unwrap_or(',')));
    let mut rows: Vec<Vec<String>> = read_rows(content, &dialect)
        .into_iter()
        .filter(|row| !row.is_empty() && !(row.len() == 1 && row[0].trim().is_empty()))
        .collect();
    if rows.is_empty() {
        return Err("no rows to chart".into());
    }
    let header = rows[0]
        .iter()
        .any(|cell| !cell.trim().is_empty() && cell.trim().parse::<f64>().is_err());
    let columns: Vec<String> = if header {
        rows.remove(0)
            .into_iter()
            .map(|c| c.trim().to_string())
            .collect()
    } else {
        let width = rows.iter().map(Vec::len).max().unwrap_or(0);
        (1..=width).map(|n| n.to_string()).collect()
    };
    if header {
        refuse_long_rows(&rows, columns.len())?;
    }
    let rows = rows
        .into_iter()
        .map(|row| {
            row.into_iter()
                .map(|cell| {
                    if cell.trim().is_empty() {
                        Value::Missing
                    } else {
                        Value::Text(cell)
                    }
                })
                .collect()
        })
        .collect();
    Ok(Data { columns, rows })
}

/// A row with more cells than the header names columns has values no
/// column holds: refused, naming the row (counted from 1, after the
/// header) rather than dropping them.
fn refuse_long_rows<T>(rows: &[Vec<T>], columns: usize) -> Result<(), String> {
    match rows.iter().position(|row| row.len() > columns) {
        Some(n) => Err(format!(
            "row {} has {} cells, but the header names {columns} columns",
            n + 1,
            rows[n].len()
        )),
        None => Ok(()),
    }
}

/// JSON: an array of records, of numbers, or of arrays (the first a header
/// when it is all strings), or an object of columns. JSON Lines: one record
/// per line.
fn parse_json(content: &str) -> Result<Data, String> {
    let document = match serde_json::from_str::<serde_json::Value>(content) {
        Ok(document) => document,
        Err(whole) => json_lines(content, Some(whole))?,
    };
    json_to_data(document)
}

/// JSON Lines: every non-blank line a record, as one array. `whole` is the
/// error from reading `content` as one document, reported instead when the
/// first line is not JSON either.
fn json_lines(
    content: &str,
    whole: Option<serde_json::Error>,
) -> Result<serde_json::Value, String> {
    use serde_json::Value as Json;
    let mut records = Vec::new();
    for (n, line) in content.lines().enumerate() {
        if line.trim().is_empty() {
            continue;
        }
        match serde_json::from_str::<Json>(line) {
            Ok(record @ Json::Object(_)) => records.push(record),
            Ok(_) => return Err(format!("line {}: expected a JSON object", n + 1)),
            Err(err) => {
                return Err(match &whole {
                    Some(whole) if records.is_empty() => format!("not valid JSON: {whole}"),
                    _ => format!("line {}: {err}", n + 1),
                });
            }
        }
    }
    Ok(Json::Array(records))
}

/// A JSON document as rows and columns.
fn json_to_data(document: serde_json::Value) -> Result<Data, String> {
    use serde_json::Value as Json;
    let value = |json: &Json| match json {
        Json::Null => Value::Missing,
        Json::Number(n) => n.as_f64().map_or(Value::Missing, Value::Number),
        Json::String(s) => Value::Text(s.clone()),
        other => Value::Text(other.to_string()),
    };
    match document {
        Json::Array(items) if items.iter().all(Json::is_object) => {
            let mut columns: Vec<String> = Vec::new();
            for item in &items {
                for key in item.as_object().into_iter().flat_map(|o| o.keys()) {
                    if !columns.contains(key) {
                        columns.push(key.clone());
                    }
                }
            }
            let rows = items
                .iter()
                .map(|item| {
                    let object = item.as_object().expect("checked above");
                    columns
                        .iter()
                        .map(|c| object.get(c).map_or(Value::Missing, value))
                        .collect()
                })
                .collect();
            Ok(Data { columns, rows })
        }
        Json::Array(items) if items.iter().all(Json::is_array) => {
            let mut rows: Vec<Vec<Value>> = items
                .iter()
                .map(|row| {
                    row.as_array()
                        .expect("checked above")
                        .iter()
                        .map(value)
                        .collect()
                })
                .collect();
            let header = items
                .first()
                .and_then(Json::as_array)
                .is_some_and(|first| !first.is_empty() && first.iter().all(Json::is_string));
            let columns: Vec<String> = if header {
                rows.remove(0).iter().map(Value::label).collect()
            } else {
                let width = rows.iter().map(Vec::len).max().unwrap_or(0);
                (1..=width).map(|n| n.to_string()).collect()
            };
            if header {
                refuse_long_rows(&rows, columns.len())?;
            }
            Ok(Data { columns, rows })
        }
        Json::Array(items) if items.iter().all(|i| !i.is_object() && !i.is_array()) => Ok(Data {
            columns: vec!["value".into()],
            rows: items.iter().map(|i| vec![value(i)]).collect(),
        }),
        Json::Object(object) if object.values().all(Json::is_array) && !object.is_empty() => {
            let columns: Vec<String> = object.keys().cloned().collect();
            let length = object
                .values()
                .filter_map(Json::as_array)
                .map(Vec::len)
                .max()
                .unwrap_or(0);
            let rows = (0..length)
                .map(|row| {
                    object
                        .values()
                        .map(|column| column.get(row).map_or(Value::Missing, value))
                        .collect()
                })
                .collect();
            Ok(Data { columns, rows })
        }
        _ => Err(
            "expected an array of records, of numbers or of rows, or an object of columns".into(),
        ),
    }
}

/// The chart `options` ask for, drawn from `data`.
fn build(data: &Data, options: &ChartOptions) -> Result<Box<dyn Renderable>, Failure> {
    let kind = options.kind.unwrap_or_default();
    let x = match &options.x {
        Some(name) => Some(data.column(name)?),
        // Labels default to the first column that is not all numbers.
        None if !kind.numeric_x() && kind != ChartKind::Spark => {
            (0..data.columns.len()).find(|&c| !data.is_numeric(c))
        }
        None => None,
    };
    let ys: Vec<usize> = if options.y.is_empty() {
        (0..data.columns.len())
            .filter(|&c| Some(c) != x && data.is_numeric(c))
            .collect()
    } else {
        options
            .y
            .iter()
            .map(|name| data.column(name))
            .collect::<Result<_, _>>()?
    };
    if ys.is_empty() {
        return Err((
            ExitClass::Data,
            format!(
                "no numeric column to chart; name one with --y (the columns are {})",
                data.columns
                    .iter()
                    .map(|c| format!("{c:?}"))
                    .collect::<Vec<_>>()
                    .join(", ")
            ),
        ));
    }
    let series: Vec<(&str, Vec<f64>)> = ys
        .iter()
        .map(|&c| Ok((data.columns[c].as_str(), data.numbers(c)?)))
        .collect::<Result<_, Failure>>()?;
    let labels: Vec<String> = (0..data.rows.len())
        .map(|row| match x {
            Some(x) => data.cell(row, x).label(),
            None => (row + 1).to_string(),
        })
        .collect();
    Ok(match kind {
        ChartKind::Spark => {
            if let [(_, values)] = series.as_slice() {
                Box::new(Sparkline::new(values.iter().copied()))
            } else {
                let mut grid = Table::grid();
                grid.add_column("");
                grid.add_column("");
                for (name, values) in &series {
                    grid.add_row_cells(vec![
                        TableCell::Text(Text::new(*name)),
                        TableCell::Renderable(std::sync::Arc::new(Sparkline::new(
                            values.iter().copied(),
                        ))),
                    ]);
                }
                Box::new(grid.padding(0, 1, 0, 0))
            }
        }
        ChartKind::Bar => {
            let mut chart = BarChart::new();
            let several = series.len() > 1;
            for (row, label) in labels.iter().enumerate() {
                for (n, (name, values)) in series.iter().enumerate() {
                    let bar = if several {
                        Bar::new(format!("{label} {name}"), values[row])
                            .style(format!("chart.series.{}", n % 5 + 1))
                    } else {
                        Bar::new(label.clone(), values[row])
                    };
                    chart = chart.push(bar);
                }
            }
            Box::new(chart)
        }
        ChartKind::Line | ChartKind::Scatter => {
            let xs: Vec<f64> = match x {
                Some(x) => data.numbers(x).map_err(|(class, err)| {
                    (
                        class,
                        format!("{err} (a line or scatter chart's --x is a position; --kind bar takes labels)"),
                    )
                })?,
                None => (1..=data.rows.len()).map(|n| n as f64).collect(),
            };
            let mut chart = LineChart::new();
            for (name, values) in &series {
                let points = xs.iter().copied().zip(values.iter().copied());
                chart = chart.series(if kind == ChartKind::Scatter {
                    Series::scatter(*name, points)
                } else {
                    Series::line(*name, points)
                });
            }
            Box::new(chart)
        }
        ChartKind::Heatmap => {
            // Each cell as wide as the widest header and a space, so every
            // header shows (the heatmap still shrinks to the width it gets).
            let header = series.iter().map(|(name, _)| cell_len(name)).max();
            let mut heatmap = Heatmap::new()
                .columns(series.iter().map(|(name, _)| *name))
                .cell_width(header.unwrap_or(1).clamp(1, 11) + 1);
            for (row, label) in labels.iter().enumerate() {
                heatmap = heatmap.row(label.clone(), series.iter().map(|(_, values)| values[row]));
            }
            Box::new(heatmap)
        }
        ChartKind::Pie => {
            let [(_, values)] = series.as_slice() else {
                return Err((
                    ExitClass::Usage,
                    format!(
                        "a pie chart draws one series; name it with --y (the numeric columns are {})",
                        series
                            .iter()
                            .map(|(name, _)| format!("{name:?}"))
                            .collect::<Vec<_>>()
                            .join(", ")
                    ),
                ));
            };
            Box::new(PieChart::from_pairs(
                labels.iter().cloned().zip(values.iter().copied()),
            ))
        }
        ChartKind::Stacked => {
            let mut chart = BarChart::new();
            for (row, label) in labels.iter().enumerate() {
                let parts = series
                    .iter()
                    .map(|(name, values)| Bar::new(*name, values[row]));
                chart = chart.group(BarGroup::new(label.clone()).stacked().bars(parts));
            }
            Box::new(chart)
        }
    })
}

#[cfg(test)]
mod tests {
    use super::*;

    fn data(content: &str, resource: &str) -> Data {
        parse(content, resource).unwrap()
    }

    #[test]
    fn a_byte_order_mark_is_dropped_before_any_parser() {
        // A URL's text keeps its BOM; JSON must still parse.
        let d = data("\u{feff}[{\"api\": 30}]", "https://example.com/data.json");
        assert_eq!(d.columns, ["api"]);
        let d = data("\u{feff}{\"api\": [1, 2]}", "https://example.com/data");
        assert_eq!(d.columns, ["api"]);
    }

    #[test]
    fn one_record_json_lines_is_a_row() {
        let d = data("{\"api\": 30}\n", "one.jsonl");
        assert_eq!(d.columns, ["api"]);
        assert_eq!(d.rows, [[Value::Number(30.0)]]);
        let d = data("{\"api\": 30}\n{\"api\": 31}\n", "two.ndjson");
        assert_eq!(d.rows.len(), 2);
    }

    #[test]
    fn csv_with_a_header_names_its_columns() {
        let d = data("name,count\napi,3\nweb,\n", "x.csv");
        assert_eq!(d.columns, ["name", "count"]);
        assert_eq!(d.rows[1][1], Value::Missing);
        assert!(d.is_numeric(1));
        assert!(!d.is_numeric(0));
    }

    #[test]
    fn csv_without_a_header_numbers_its_columns() {
        let d = data("1,2\n3,4\n", "-");
        assert_eq!(d.columns, ["1", "2"]);
        assert_eq!(d.column("2").unwrap(), 1);
    }

    #[test]
    fn bare_numbers_are_one_series() {
        let d = data("1\n2 3\n", "-");
        assert_eq!(d.columns, ["value"]);
        assert_eq!(d.rows.len(), 3);
    }

    #[test]
    fn json_shapes() {
        let records = data(r#"[{"a":1,"b":2},{"a":3,"c":null}]"#, "-");
        assert_eq!(records.columns, ["a", "b", "c"]);
        assert_eq!(records.rows[1][1], Value::Missing);
        let numbers = data("[1, 2, null]", "-");
        assert_eq!(numbers.columns, ["value"]);
        let rows = data(r#"[["x","y"],[1,2]]"#, "-");
        assert_eq!(rows.columns, ["x", "y"]);
        let columns = data(r#"{"a":[1,2],"b":[3]}"#, "-");
        assert_eq!(columns.rows.len(), 2);
        assert_eq!(columns.rows[1][1], Value::Missing);
        let lines = data("{\"a\":1}\n{\"a\":2}\n", "x.jsonl");
        assert_eq!(lines.rows.len(), 2);
    }

    #[test]
    fn errors_name_the_row_and_the_column() {
        let d = data("t,v\n1,2\n2,oops\n", "x.csv");
        let err = d.numbers(1).unwrap_err();
        assert_eq!(err.1, "row 2, column \"v\": \"oops\" is not a number");
        let err = d.column("w").unwrap_err();
        assert_eq!(err.1, "no column \"w\"; the columns are \"t\", \"v\"");
        assert!(parse("", "-").is_err());
        assert!(parse("{\"a\": 1}", "-").is_err());
    }

    #[test]
    fn labels_write_numbers_back_as_read() {
        assert_eq!(Value::Number(2024.0).label(), "2024");
        assert_eq!(Value::Number(1.25).label(), "1.25");
    }
}