rs-rich-cli 0.0.21

//! `rich profile`: a profile of CSV, TSV or JSON Lines (0.0.16 workstream 3,
//! #343-#346; nested JSON Lines paths in 0.0.19, #705). Not upstream: a
//! binary-boundary convenience that composes `rs-rich-data`'s `profile`; see
//! `docs/PORTING.md`.
//!
//! This file is command routing: choosing the reader (by extension, else by
//! the first character, as `rich chart` chooses), streaming a file or stdin
//! into it so a long input is sampled rather than read into memory, and
//! `--report json`. The profile itself is `rich_data::profile`'s: JSON Lines
//! are profiled path by path (`JsonProfiler`) unless `--flat` asks for the
//! rows of top-level keys.
use super::*;

use std::io::BufReader;

use rich_data::csv::{CsvReader, Header};
use rich_data::profile::{Profile, ProfileOptions as DataProfileOptions};
use rich_data::{jsonl, RowSource};

/// The options of `rich profile`.
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub(crate) struct ProfileOptions {
    /// `--sample N`: the most rows kept for types and statistics.
    sample: Option<usize>,
    /// `--columns a,b`: only these columns, in this order.
    columns: Option<Vec<String>>,
    /// `--top N`: the most common values shown per categorical column.
    top: Option<usize>,
    /// `--depth N`: the deepest JSON path profiled. Parsed by `rich deps`'s
    /// options, which own the flag, and handed over by
    /// [`set_depth`](Self::set_depth).
    depth: Option<usize>,
    /// `--max-paths N`: the most JSON paths profiled.
    max_paths: Option<usize>,
    /// `--flat`: JSON Lines as rows of their top-level keys.
    flat: bool,
}

fn count(flag: &str, value: Option<&String>) -> Result<usize, String> {
    value
        .and_then(|v| v.trim().replace('_', "").parse::<usize>().ok())
        .filter(|n| *n > 0)
        .ok_or_else(|| format!("{flag} requires a positive integer"))
}

impl ProfileOptions {
    /// 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 {
            "--sample" => self.sample = Some(count(arg, rest.next())?),
            "--top" => self.top = Some(count(arg, rest.next())?),
            "--max-paths" => self.max_paths = Some(count(arg, rest.next())?),
            "--flat" => self.flat = true,
            "--columns" => {
                let value = rest.next().ok_or("--columns requires COLUMN[,COLUMN...]")?;
                let columns: Vec<String> = value
                    .split(',')
                    .map(|c| c.trim().to_string())
                    .filter(|c| !c.is_empty())
                    .collect();
                if columns.is_empty() {
                    return Err("--columns requires COLUMN[,COLUMN...]".into());
                }
                self.columns.get_or_insert_with(Vec::new).extend(columns);
            }
            _ => 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])> {
        [
            ("--sample", self.sample.is_some()),
            ("--columns", self.columns.is_some()),
            ("--top", self.top.is_some()),
            ("--max-paths", self.max_paths.is_some()),
            ("--flat", self.flat),
        ]
        .into_iter()
        .filter(|(_, given)| *given)
        .map(|(flag, _)| (flag, &["profile"][..]))
        .collect()
    }

    /// `--depth N` for `rich profile`, which `rich deps`'s options parsed.
    pub(crate) fn set_depth(&mut self, depth: usize) -> Result<(), String> {
        if depth == 0 {
            return Err("--depth requires a positive integer".into());
        }
        self.depth = Some(depth);
        Ok(())
    }

    /// The JSON Lines options given (`--depth`, `--max-paths`, `--flat`),
    /// which CSV has no use for.
    fn json_only(&self) -> Vec<&'static str> {
        [
            ("--depth", self.depth.is_some()),
            ("--max-paths", self.max_paths.is_some()),
            ("--flat", self.flat),
        ]
        .into_iter()
        .filter_map(|(flag, given)| given.then_some(flag))
        .collect()
    }

    fn data_options(&self) -> DataProfileOptions {
        let defaults = DataProfileOptions::default();
        DataProfileOptions {
            sample: self.sample.unwrap_or(defaults.sample),
            top: self.top.unwrap_or(defaults.top),
            columns: self.columns.clone(),
            depth: self.depth.unwrap_or(defaults.depth),
            max_paths: self.max_paths.unwrap_or(defaults.max_paths),
            ..defaults
        }
    }
}

type Failure = (ExitClass, String);

/// How the input is read.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Format {
    /// CSV or TSV, sniffed; the delimiter to fall back to.
    Csv(Option<char>),
    JsonLines,
}

/// The format by the resource's extension, else `None`.
fn format_by_name(resource: &str) -> Option<Format> {
    let path = resource.split(['?', '#']).next().unwrap_or(resource);
    let extension = path
        .rsplit(['/', '\\'])
        .next()
        .and_then(|name| name.rsplit_once('.'))
        .map(|(_, ext)| ext.to_ascii_lowercase());
    match extension.as_deref() {
        Some("jsonl" | "ndjson") => Some(Format::JsonLines),
        Some("csv") => Some(Format::Csv(Some(','))),
        Some("tsv" | "tab") => Some(Format::Csv(Some('\t'))),
        _ => None,
    }
}

/// `rich profile [FILE]`: stream the input into a profile.
pub(crate) fn profile(cli: &Cli) -> Result<Profile, Failure> {
    let resource = cli.resource.as_deref().unwrap_or("-");
    let name = if resource == "-" { "<stdin>" } else { resource };
    let input_error =
        |err: std::io::Error| (ExitClass::Input, format!("cannot read {name}: {err}"));
    // Files and stdin stream; a URL's body (already bounded by the fetch
    // limit) and text in another `--encoding` are decoded whole first.
    let mut reader: Box<dyn BufRead> = if is_url(resource) {
        let (text, _) =
            fetch_url(resource, cli.extensions.encoding).map_err(|err| (ExitClass::Input, err))?;
        Box::new(std::io::Cursor::new(text.into_bytes()))
    } else if cli.extensions.encoding.is_some() {
        let text = read_resource(Some(resource), cli.extensions.encoding).map_err(input_error)?;
        Box::new(std::io::Cursor::new(text.into_bytes()))
    } else if resource == "-" {
        if std::io::stdin().is_terminal() {
            let eof = if cfg!(windows) {
                "Ctrl-Z then Enter"
            } else {
                "Ctrl-D"
            };
            eprintln!("rich: reading stdin; finish input with {eof}");
        }
        Box::new(BufReader::new(std::io::stdin()))
    } else {
        let path = fs_path(resource);
        if path.is_dir() {
            return Err((
                ExitClass::Input,
                format!("cannot read {name}: is a directory, not a file"),
            ));
        }
        Box::new(BufReader::new(
            std::fs::File::open(path).map_err(input_error)?,
        ))
    };
    let format = match format_by_name(resource) {
        Some(format) => format,
        None => {
            let (first, prefix) = first_byte(&mut reader).map_err(input_error)?;
            // The bytes read to find it go back in front, unchanged.
            reader = Box::new(BufReader::new(std::io::Cursor::new(prefix).chain(reader)));
            match first {
                Some(b'{' | b'[') => Format::JsonLines,
                _ => Format::Csv(None),
            }
        }
    };
    let data_error = |err: rich_data::DataError| (ExitClass::Data, format!("{name}: {err}"));
    let options = cli.profile.data_options();
    if matches!(format, Format::Csv(_)) {
        let flags = cli.profile.json_only();
        if !flags.is_empty() {
            return Err((
                ExitClass::Usage,
                format!(
                    "{} only {} JSON Lines; {name} is read as CSV",
                    flags.join(", "),
                    if flags.len() == 1 {
                        "applies to"
                    } else {
                        "apply to"
                    }
                ),
            ));
        }
    }
    let profile = match format {
        // Path by path: nested objects and arrays, presence and type mixes;
        // a line that is not JSON is counted, not fatal.
        Format::JsonLines if !cli.profile.flat => {
            let profile = Profile::from_json_lines(reader, options).map_err(data_error)?;
            if profile.rows() == 0 && profile.malformed() > 0 {
                return Err((
                    ExitClass::Data,
                    format!(
                        "{name}: no records to profile ({})",
                        profile
                            .notes()
                            .last()
                            .map(String::as_str)
                            .unwrap_or_default()
                    ),
                ));
            }
            profile
        }
        Format::JsonLines => {
            let mut source = jsonl::source(reader).map_err(data_error)?;
            let mut profile = run(name, &mut source, options)?;
            let mut late: Vec<String> = source.unknown_keys().map(controls::shown).collect();
            if source.more_unknown_keys() {
                late.push("…".into());
            }
            if !late.is_empty() {
                profile.note(format!(
                    "keys first seen after record {} are not profiled: {}",
                    jsonl::DEFAULT_SAMPLE,
                    late.join(", ")
                ));
            }
            profile
        }
        Format::Csv(fallback) => {
            // As `rich chart`: a comma when the sniffer finds no delimiter
            // (one column, or no input at all).
            let mut source = CsvReader::new()
                .header_mode(Header::UnlessNumeric)
                .fallback(fallback.unwrap_or(','))
                .source(reader)
                .map_err(data_error)?;
            let mut profile = run(name, &mut source, options)?;
            let longer = source.longer_rows();
            if longer > 0 {
                profile.note(format!(
                    "{longer} {} more cells than the header names; the extra cells are not \
                     profiled",
                    if longer == 1 { "row has" } else { "rows have" }
                ));
            }
            profile
        }
    };
    if profile.columns().is_empty() && cli.profile.columns.is_none() {
        return Err((ExitClass::Data, format!("{name}: no rows to profile")));
    }
    // `--style-rule` and `--rules`, by the table's headings.
    let rules = cli.rules.rules();
    if !rules.is_empty() {
        cli.rules.warn_unresolved(
            profile.table_headers(),
            cli.report_format == ReportFormat::Json,
        );
    }
    Ok(profile
        .with_name(controls::shown(name))
        .style_rules(rules.clone()))
}

/// The most leading whitespace read while looking for the first character
/// (`rich_data::csv::MAX_RECORD`, 64 MiB); input that is only whitespace up
/// to it is read as CSV.
const MAX_LEADING_WHITESPACE: usize = rich_data::csv::MAX_RECORD;

/// The first byte of `reader` that is not whitespace or a leading
/// byte-order mark, read across as many buffer refills as the whitespace
/// takes (up to [`MAX_LEADING_WHITESPACE`]), with every byte read on the
/// way so the caller can put them back.
fn first_byte(reader: &mut dyn BufRead) -> std::io::Result<(Option<u8>, Vec<u8>)> {
    const BOM: &[u8] = "\u{feff}".as_bytes();
    let mut prefix: Vec<u8> = Vec::new();
    // Bytes already known to be whitespace (or the mark).
    let mut scanned = 0;
    loop {
        let chunk = reader.fill_buf()?;
        if chunk.is_empty() {
            return Ok((None, prefix));
        }
        let len = chunk.len();
        prefix.extend_from_slice(chunk);
        reader.consume(len);
        if scanned == 0 {
            // A byte-order mark is only skipped once all of it is read.
            if prefix.starts_with(BOM) {
                scanned = BOM.len();
            } else if BOM.starts_with(&prefix) {
                continue;
            }
        }
        if let Some(&b) = prefix[scanned..].iter().find(|b| !b.is_ascii_whitespace()) {
            return Ok((Some(b), prefix));
        }
        scanned = prefix.len();
        if prefix.len() > MAX_LEADING_WHITESPACE {
            return Ok((None, prefix));
        }
    }
}

/// Profile `source`, a column `--columns` names that is not there and a row
/// that cannot be read being data errors.
fn run(
    name: &str,
    source: &mut impl RowSource,
    options: DataProfileOptions,
) -> Result<Profile, Failure> {
    let mut profiler = rich_data::profile::Profiler::new(source.columns(), options)
        .map_err(|err| (ExitClass::Data, format!("{name}: {err}")))?;
    while let Some(row) = source.next_row() {
        profiler.push(row.map_err(|err| (ExitClass::Data, format!("{name}: {err}")))?);
    }
    Ok(profiler.finish())
}

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

    #[test]
    fn the_format_comes_from_the_extension() {
        assert_eq!(format_by_name("a.CSV"), Some(Format::Csv(Some(','))));
        assert_eq!(format_by_name("dir.x/a.tsv"), Some(Format::Csv(Some('\t'))));
        assert_eq!(
            format_by_name("https://h/e.ndjson?x=1"),
            Some(Format::JsonLines)
        );
        assert_eq!(format_by_name("data"), None);
        assert_eq!(format_by_name("-"), None);
    }

    #[test]
    fn the_first_byte_is_found_past_any_buffer_of_whitespace() {
        let text = format!("\u{feff}{}{{\"a\": 1}}\n", " \n\t".repeat(10_000));
        // A 4-byte buffer splits the mark and needs thousands of refills.
        let mut reader = BufReader::with_capacity(4, text.as_bytes());
        let (first, prefix) = first_byte(&mut reader).unwrap();
        assert_eq!(first, Some(b'{'));
        let mut rest = prefix;
        reader.read_to_end(&mut rest).unwrap();
        assert_eq!(rest, text.as_bytes());
        let (first, prefix) = first_byte(&mut BufReader::new(&b"\n \n"[..])).unwrap();
        assert_eq!((first, prefix.as_slice()), (None, &b"\n \n"[..]));
        let (first, _) = first_byte(&mut BufReader::new(&b"\xef\xbb"[..])).unwrap();
        assert_eq!(first, None);
    }

    #[test]
    fn counts_are_positive() {
        let mut options = ProfileOptions::default();
        let args = ["10_000".to_string()];
        assert!(options.parse_option("--sample", &mut args.iter()).unwrap());
        assert_eq!(options.sample, Some(10_000));
        let zero = ["0".to_string()];
        assert!(options.parse_option("--top", &mut zero.iter()).is_err());
        let empty = [",".to_string()];
        assert!(options
            .parse_option("--columns", &mut empty.iter())
            .is_err());
    }
}