Skip to content

Commit

Permalink
sqlp: automatically use tab delimiters when output format is CSV an…
Browse files Browse the repository at this point in the history
…d file extension is TSV or TAB (like the rest of qsv)
  • Loading branch information
jqnatividad committed Jan 3, 2024
1 parent 7ac0602 commit c97048c
Show file tree
Hide file tree
Showing 2 changed files with 46 additions and 2 deletions.
17 changes: 15 additions & 2 deletions src/cmd/sqlp.rs
Original file line number Diff line number Diff line change
Expand Up @@ -240,7 +240,7 @@ struct Args {
flag_quiet: bool,
}

#[derive(Default, Clone)]
#[derive(Default, Clone, PartialEq)]
enum OutputMode {
#[default]
Csv,
Expand All @@ -258,7 +258,7 @@ impl OutputMode {
&self,
query: &str,
ctx: &mut SQLContext,
delim: u8,
mut delim: u8,
args: Args,
) -> CliResult<(usize, usize)> {
let mut df = DataFrame::default();
Expand All @@ -275,6 +275,19 @@ impl OutputMode {
let w = match args.flag_output {
Some(x) => {
let path = Path::new(&x);
// if the output file ends with ".tsv" or ".tab", we use tab as the delimiter
delim = if *self == OutputMode::Csv
&& path
.extension()
.map_or(false, |ext| ext.eq_ignore_ascii_case("tsv"))
|| path
.extension()
.map_or(false, |ext| ext.eq_ignore_ascii_case("tab"))
{
b'\t'
} else {
delim
};
Box::new(File::create(path)?) as Box<dyn Write>
},
None => Box::new(io::stdout()) as Box<dyn Write>,
Expand Down
31 changes: 31 additions & 0 deletions tests/test_sqlp.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1149,3 +1149,34 @@ fn sqlp_sql_from_subquery() {

assert_eq!(got, expected);
}

#[test]
fn sqlp_sql_tsv() {
let wrk = Workdir::new("sqlp_sql_tsv");
wrk.create(
"test.csv",
vec![
svec!["idx", "val"],
svec!["0", "ABC"],
svec!["1", "abc"],
svec!["2", "000"],
svec!["3", "A0C"],
svec!["4", "a0c"],
],
);

let output_file = wrk.path("output.tsv").to_string_lossy().to_string();

let mut cmd = wrk.command("sqlp");
cmd.arg("test.csv")
.arg("SELECT * FROM test")
.args(["--output", &output_file]);

wrk.assert_success(&mut cmd);

let got = wrk.read_to_string(&output_file);

let expected = "idx\tval\n0\tABC\n1\tabc\n2\t000\n3\tA0C\n4\ta0c\n";

assert_eq!(got, expected);
}

0 comments on commit c97048c

Please sign in to comment.