aboutsummaryrefslogtreecommitdiffstats
path: root/src/bin/aft-head.rs
diff options
context:
space:
mode:
Diffstat (limited to 'src/bin/aft-head.rs')
-rw-r--r--src/bin/aft-head.rs69
1 files changed, 66 insertions, 3 deletions
diff --git a/src/bin/aft-head.rs b/src/bin/aft-head.rs
index b8c848c..3fecfa2 100644
--- a/src/bin/aft-head.rs
+++ b/src/bin/aft-head.rs
@@ -1,8 +1,71 @@
-
/*
* Takes the first n lines of input, passes to stdout.
*/
-fn main() {
- println!("Hello, world2!");
+use std::error::Error;
+use std::io::{self, Read, Write, BufRead};
+use structopt::StructOpt;
+use std::path::PathBuf;
+
+#[derive(StructOpt)]
+#[structopt(rename_all = "kebab-case", about = "take first N rows of input (AFT-aware)" )]
+struct Opt {
+ #[structopt(long = "--lines", short = "-n", default_value = "10")]
+ lines: u64,
+
+ #[structopt(parse(from_os_str))]
+ input: Option<PathBuf>,
+}
+
+fn main() -> Result<(), Box<dyn Error>> {
+
+ let opt = Opt::from_args();
+
+ // treat "-" as "use stdin"
+ let input = match opt.input {
+ Some(s) if s.to_string_lossy() == "-" => None,
+ other => other,
+ };
+ let stdin = io::stdin();
+ let stdin = stdin.lock();
+ let mut input: Box<dyn BufRead> = match input {
+ None => {
+ Box::new(io::BufReader::new(stdin))
+ },
+ Some(path) => {
+ let f = std::fs::File::open(path)?;
+ Box::new(io::BufReader::new(f))
+ },
+ };
+ let mut stdout = io::stdout();
+
+ // read first byte of input to check if we are getting AFT; if not just pass through
+ let mut first_byte = [0; 1];
+ let got = input.read(&mut first_byte)?;
+ if got != 1 {
+ panic!("couldn't read a byte from input");
+ };
+
+ let is_aft = first_byte[0] == 0x01;
+ if !is_aft {
+ // simple implementation of lines
+ stdout.write(&first_byte)?;
+ for l in input.lines().take(opt.lines as usize) {
+ writeln!(stdout, "{}", l?)?;
+ }
+ return Ok(())
+ };
+
+ let mut header_bytes = vec![];
+ // TODO: sanity limit on length
+ input.read_until(0x02, &mut header_bytes)?;
+ if header_bytes[header_bytes.len()-1] != 0x02 {
+ panic!("expected an AFT header");
+ };
+ stdout.write(&first_byte)?;
+ stdout.write(&header_bytes)?;
+ for l in input.lines().take(opt.lines as usize) {
+ writeln!(stdout, "{}", l?)?;
+ }
+ Ok(())
}