From 3f2153e08a6b27ca96b2c0cc9a728da817f310be Mon Sep 17 00:00:00 2001 From: Linlin Sun Date: Fri, 10 Mar 2023 15:03:49 -0800 Subject: [PATCH] Adds 'head' option to allow user get the first specific number of ion values of the input file. --- src/bin/ion/commands/beta/head.rs | 54 +++++++++++++++++++++++++++++++ src/bin/ion/commands/beta/mod.rs | 4 ++- src/bin/ion/commands/dump.rs | 23 ++++++++----- src/bin/ion/commands/mod.rs | 1 - tests/cli.rs | 16 +++++++++ tests/testData.ion | 16 +++++++++ 6 files changed, 104 insertions(+), 10 deletions(-) create mode 100644 src/bin/ion/commands/beta/head.rs create mode 100644 tests/testData.ion diff --git a/src/bin/ion/commands/beta/head.rs b/src/bin/ion/commands/beta/head.rs new file mode 100644 index 0000000..3abd732 --- /dev/null +++ b/src/bin/ion/commands/beta/head.rs @@ -0,0 +1,54 @@ +use std::io::{stdout, Write}; +use clap::{Arg, ArgMatches, Command}; +use ion_rs::*; +use std::fs::File; +use anyhow::{Context, Result}; +use crate::commands::dump; + +pub fn app() -> Command { + Command::new("head") + .about("Prints the specified number of top-level values in the input stream.") + .arg( + Arg::new("number") + .long("number") + .short('n') + .allow_negative_numbers(false) + .help("Specifies the number of output top-level values.") + ) + .arg( + Arg::new("format") + .long("format") + .short('f') + .default_value("lines") + .value_parser(["binary", "text", "pretty", "lines"]) + .help("Output format"), + ) + .arg( + Arg::new("input") + .required(true) + .help("Input file") + ) +} + +pub fn run(_command_name: &str, matches: &ArgMatches) -> Result<()> { + let input_file = matches + .get_one::("input") + .unwrap(); + let number: i32 = matches + .get_one::("number") + .unwrap() + .parse() + .unwrap(); + let format = matches + .get_one::("format") + .expect("`format` did not have a value"); + let file = File::open(input_file) + .with_context(|| format!("Could not open file '{}'", input_file))?; + if number == 0 { + return Ok(()); + } + let mut reader = ReaderBuilder::new().build(file)?; + let mut output: Box = Box::new(stdout().lock()); + dump::write_all_in_format(&mut reader, &mut output, format, number).expect("Could not process the input stream."); + Ok(()) +} \ No newline at end of file diff --git a/src/bin/ion/commands/beta/mod.rs b/src/bin/ion/commands/beta/mod.rs index 4597aa9..40fb077 100644 --- a/src/bin/ion/commands/beta/mod.rs +++ b/src/bin/ion/commands/beta/mod.rs @@ -2,7 +2,7 @@ pub mod count; pub mod inspect; pub mod primitive; pub mod schema; - +pub mod head; use crate::commands::CommandRunner; use anyhow::Result; use clap::{ArgMatches, Command}; @@ -17,6 +17,7 @@ pub fn beta_subcommands() -> Vec { inspect::app(), primitive::app(), schema::app(), + head::app(), ] } @@ -26,6 +27,7 @@ pub fn runner_for_beta_subcommand(command_name: &str) -> Option { "inspect" => inspect::run, "primitive" => primitive::run, "schema" => schema::run, + "head" => head::run, _ => return None, }; Some(runner) diff --git a/src/bin/ion/commands/dump.rs b/src/bin/ion/commands/dump.rs index e8e23df..48d19d7 100644 --- a/src/bin/ion/commands/dump.rs +++ b/src/bin/ion/commands/dump.rs @@ -33,6 +33,7 @@ pub fn app() -> Command { } pub fn run(_command_name: &str, matches: &ArgMatches) -> Result<()> { + const HEAD_NUMBER:i32 = -1; // --format pretty|text|lines|binary // `clap` validates the specified format and provides a default otherwise. let format = matches @@ -58,12 +59,12 @@ pub fn run(_command_name: &str, matches: &ArgMatches) -> Result<()> { let file = File::open(input_file) .with_context(|| format!("Could not open file '{}'", input_file))?; let mut reader = ReaderBuilder::new().build(file)?; - write_all_in_format(&mut reader, &mut output, format)?; + write_all_in_format(&mut reader, &mut output, format, HEAD_NUMBER)?; } } else { let input: StdinLock = stdin().lock(); let mut reader = ReaderBuilder::new().build(input)?; - write_all_in_format(&mut reader, &mut output, format)?; + write_all_in_format(&mut reader, &mut output, format, HEAD_NUMBER)?; } output.flush()?; @@ -72,27 +73,28 @@ pub fn run(_command_name: &str, matches: &ArgMatches) -> Result<()> { /// Constructs the appropriate writer for the given format, then writes all values found in the /// Reader to the new Writer. -fn write_all_in_format( +pub(crate) fn write_all_in_format( reader: &mut Reader, output: &mut Box, format: &str, + head_number: i32 ) -> IonResult<()> { match format { "pretty" => { let mut writer = TextWriterBuilder::pretty().build(output)?; - write_all_values(reader, &mut writer) + write_all_values(reader, &mut writer, head_number) } "text" => { let mut writer = TextWriterBuilder::default().build(output)?; - write_all_values(reader, &mut writer) + write_all_values(reader, &mut writer, head_number) } "lines" => { let mut writer = TextWriterBuilder::lines().build(output)?; - write_all_values(reader, &mut writer) + write_all_values(reader, &mut writer, head_number) } "binary" => { let mut writer = BinaryWriterBuilder::new().build(output)?; - write_all_values(reader, &mut writer) + write_all_values(reader, &mut writer, head_number) } unrecognized => unreachable!( "'format' was '{}' instead of 'pretty', 'text', 'lines', or 'binary'", @@ -102,10 +104,11 @@ fn write_all_in_format( } /// Writes each value encountered in the Reader to the provided IonWriter. -fn write_all_values(reader: &mut Reader, writer: &mut W) -> IonResult<()> { +fn write_all_values(reader: &mut Reader, writer: &mut W, head_number: i32) -> IonResult<()> { const FLUSH_EVERY_N: usize = 100; let mut values_since_flush: usize = 0; let mut annotations = vec![]; + let mut index = 0; loop { match reader.next()? { StreamItem::Value(ion_type) | StreamItem::Null(ion_type) => { @@ -168,12 +171,16 @@ fn write_all_values(reader: &mut Reader, writer: &mut W) -> IonRes StreamItem::Nothing => break, } if reader.depth() == 0 { + index += 1; values_since_flush += 1; if values_since_flush == FLUSH_EVERY_N { writer.flush()?; values_since_flush = 0; } } + if head_number != -1 && index == head_number { + break; + } } writer.flush()?; Ok(()) diff --git a/src/bin/ion/commands/mod.rs b/src/bin/ion/commands/mod.rs index 9751239..ae4cbd6 100644 --- a/src/bin/ion/commands/mod.rs +++ b/src/bin/ion/commands/mod.rs @@ -1,6 +1,5 @@ use anyhow::Result; use clap::{ArgMatches, Command}; - pub mod beta; pub mod dump; diff --git a/tests/cli.rs b/tests/cli.rs index 1f397b8..ef70b56 100644 --- a/tests/cli.rs +++ b/tests/cli.rs @@ -144,3 +144,19 @@ fn run_it>( Ok(()) } + +#[rstest( +number, expected_output, +case(0, ""), +case(2, "{foo: bar, abc: [123, 456]}\n{foo: baz, abc: [420d-1, 4.3e1]}"), +)] +///Calls ion-cli beta head with different requested number. Pass the test if the return value equals to the expected value. +fn test_head(number: i32, expected_output: &str) -> Result<()> { + let mut cmd = Command::cargo_bin("ion")?; + cmd.args(&["beta", "head", "--number", &number.to_string(), "--format", "lines", "./tests/testData.ion"]); + let command_assert = cmd.assert(); + let output = command_assert.get_output(); + let stdout = String::from_utf8_lossy(&output.stdout); + assert_eq!(stdout.trim_end(), expected_output); + Ok(()) +} \ No newline at end of file diff --git a/tests/testData.ion b/tests/testData.ion new file mode 100644 index 0000000..b1b11b5 --- /dev/null +++ b/tests/testData.ion @@ -0,0 +1,16 @@ +{ + foo: bar, + abc: [123, 456] +} +{ + foo: baz, + abc: [42.0, 43e0] +} +{ + foo: bar, + test: data +} +{ + foo: baz, + type: struct +}