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

@@ -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::<String>("input")
+        .unwrap();
+    let number: i32 = matches
+        .get_one::<String>("number")
+        .unwrap()
+        .parse()
+        .unwrap();
+    let format = matches
+        .get_one::<String>("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<dyn Write> = Box::new(stdout().lock());
+    dump::write_all_in_format(&mut reader, &mut output, format, number).expect("Could not process the input stream.");
+    Ok(())
+}

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<Command> {
         inspect::app(),
         primitive::app(),
         schema::app(),
+        head::app(),
     ]
 }
 
@@ -26,6 +27,7 @@ pub fn runner_for_beta_subcommand(command_name: &str) -> Option<CommandRunner> {
         "inspect" => inspect::run,
         "primitive" => primitive::run,
         "schema" => schema::run,
+        "head" => head::run,
         _ => return None,
     };
     Some(runner)

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<dyn Write>,
     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<W: IonWriter>(reader: &mut Reader, writer: &mut W) -> IonResult<()> {
+fn write_all_values<W: IonWriter>(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<W: IonWriter>(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(())

src/bin/ion/commands/mod.rs

@@ -1,6 +1,5 @@
 use anyhow::Result;
 use clap::{ArgMatches, Command};
-
 pub mod beta;
 pub mod dump;
 

tests/cli.rs

@@ -144,3 +144,19 @@ fn run_it<S: AsRef<str>>(
 
     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(())
+}

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
+}