aboutsummaryrefslogtreecommitdiffstats
path: root/src/bin/aft.rs
blob: 25aa0f756d0aed357436ec10c86dc32112259db9 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105

/*
 * Default behavior:
 *
 *   aft < file.aft
 *      pretty-prints as a table to stdout
 *
 *   ls | aft | aft-sort
 *      detects non-aft stdin and converts to aft on stdout
 */

use std::str::FromStr;
use std::io::{self, Read, Write, BufRead};
use std::error::Error;
use colored::*;

#[derive(Debug)]
enum AftFieldType {
    Str,
    Int,
    Float,
    Bool,
    Array,
    Null,
    Other(String),
}

impl FromStr for AftFieldType {
    type Err = Box<dyn Error>;

    fn from_str(raw: &str) -> Result<AftFieldType, Box<dyn Error>> {
        match raw {
            "" => Ok(AftFieldType::Null),
            "str" => Ok(AftFieldType::Str),
            "int" | "integer" => Ok(AftFieldType::Int),
            "float" => Ok(AftFieldType::Float),
            "bool" | "boolean" => Ok(AftFieldType::Bool),
            "array" => Ok(AftFieldType::Array),
            other => Ok(AftFieldType::Other(other.to_string())),

        }
    }
}

#[derive(Debug)]
struct AftHeader {
    col_names: Vec<String>,
    col_types: Option<Vec<AftFieldType>>,
}

fn main() -> Result<(), Box<dyn Error>> {

    let mut stdout = io::stdout();
    let stdin = io::stdin();
    let mut stdin = stdin.lock();

    // if we aren't connected to terminal, just pass through
    if !atty::is(atty::Stream::Stdout) {
        io::copy(&mut stdin, &mut stdout)?;
        return Ok(())
    }

    // 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 = stdin.read(&mut first_byte)?;
    if got != 1 {
        panic!("couldn't read a byte from stdin");
    };
    if first_byte[0] != 0x01 {
        stdout.write(&first_byte)?;
        io::copy(&mut stdin, &mut stdout)?;
        return Ok(())
    }

    let mut stdin = io::BufReader::new(stdin);
    let mut header_bytes = vec![];
    // TODO: sanity limit on length
    stdin.read_until(0x02, &mut header_bytes)?;
    if header_bytes[header_bytes.len()-1] != 0x02 {
        panic!("expected an AFT header");
    };
    let header_string = String::from_utf8_lossy(&header_bytes);
    let header_rows: Vec<&str> = header_string.splitn(2, "\n").collect();
    let header = match header_rows.len() {
        0 => panic!("expected a header"),
        1 => AftHeader {
            col_names: header_rows[0].split("\x1E").map(|v| v.to_string()).collect(),
            col_types: None,
        },
        _ => AftHeader {
            col_names: header_rows[0].split("\x1E").map(|v| v.to_string()).collect(),
            col_types: Some(header_rows[1].split("\x1E").map(|v| AftFieldType::from_str(v)).collect::<Result<Vec<AftFieldType>, Box<dyn Error>>>()?),
        },
    };

    let mut tw = tabwriter::TabWriter::new(stdout);
    writeln!(tw, "{}", header.col_names.join("\t").bold())?;
    for line in stdin.lines() {
        let line = line?;
        writeln!(tw, "{}", line.replace("\x1E", "\t").replace("\x1D", ""))?;
    }
    tw.flush()?;

    Ok(())
}