Skip to content

Commit

Permalink
Support reading from stdin in join
Browse files Browse the repository at this point in the history
  • Loading branch information
pjsier committed Oct 4, 2019
1 parent c0d2666 commit 7e9f503
Show file tree
Hide file tree
Showing 2 changed files with 22 additions and 5 deletions.
9 changes: 4 additions & 5 deletions src/cmd/join.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,5 @@
use std::collections::hash_map::{HashMap, Entry};
use std::fmt;
use std::fs;
use std::io;
use std::iter::repeat;
use std::str;
Expand All @@ -9,7 +8,7 @@ use byteorder::{WriteBytesExt, BigEndian};
use csv;

use CliResult;
use config::{Config, Delimiter};
use config::{Config, Delimiter, SeekRead};
use index::Indexed;
use select::{SelectColumns, Selection};
use util;
Expand Down Expand Up @@ -278,7 +277,7 @@ impl<R: io::Read + io::Seek, W: io::Write> IoState<R, W> {

impl Args {
fn new_io_state(&self)
-> CliResult<IoState<fs::File, Box<io::Write+'static>>> {
-> CliResult<IoState<Box<SeekRead+'static>, Box<io::Write+'static>>> {
let rconf1 = Config::new(&Some(self.arg_input1.clone()))
.delimiter(self.flag_delimiter)
.no_headers(self.flag_no_headers)
Expand All @@ -288,8 +287,8 @@ impl Args {
.no_headers(self.flag_no_headers)
.select(self.arg_columns2.clone());

let mut rdr1 = rconf1.reader_file()?;
let mut rdr2 = rconf2.reader_file()?;
let mut rdr1 = rconf1.reader_file_stdin()?;
let mut rdr2 = rconf2.reader_file_stdin()?;
let (sel1, sel2) = self.get_selections(
&rconf1, &mut rdr1, &rconf2, &mut rdr2)?;
Ok(IoState {
Expand Down
18 changes: 18 additions & 0 deletions src/config.rs
Original file line number Diff line number Diff line change
Expand Up @@ -70,6 +70,10 @@ pub struct Config {
quoting: bool,
}

// Empty trait as an alias for Seek and Read that avoids auto trait errors
pub trait SeekRead: io::Seek + io::Read {}
impl<T: io::Seek + io::Read> SeekRead for T {}

impl Config {
pub fn new(path: &Option<String>) -> Config {
let (path, delim) = match *path {
Expand Down Expand Up @@ -212,6 +216,20 @@ impl Config {
}
}

pub fn reader_file_stdin(&self) -> io::Result<csv::Reader<Box<SeekRead+'static>>> {
Ok(match self.path {
None => {
// Create a buffer in memory when stdin needs to be indexed
let mut buffer: Vec<u8> = Vec::new();
io::stdin().lock().read_to_end(&mut buffer)?;
self.from_reader(Box::new(io::Cursor::new(buffer)))
},
Some(ref p) => {
self.from_reader(Box::new(fs::File::open(p).unwrap()))
},
})
}

pub fn index_files(&self)
-> io::Result<Option<(csv::Reader<fs::File>, fs::File)>> {
let (csv_file, idx_file) = match (&self.path, &self.idx_path) {
Expand Down

0 comments on commit 7e9f503

Please sign in to comment.