Accept queries from standard input when not given as argument

This commit is contained in:
Kerollmops 2020-06-20 12:01:15 +02:00
parent d6705d5529
commit c84012d655
No known key found for this signature in database
GPG Key ID: 92ADA4E935E71FA4

View File

@ -1,4 +1,5 @@
use std::io::{self, Write}; use std::io::{self, Write, BufRead};
use std::iter::once;
use std::path::PathBuf; use std::path::PathBuf;
use std::time::Instant; use std::time::Instant;
@ -19,7 +20,7 @@ struct Opt {
database: PathBuf, database: PathBuf,
/// The query string to search for (doesn't support prefix search yet). /// The query string to search for (doesn't support prefix search yet).
query: String, query: Option<String>,
} }
fn main() -> anyhow::Result<()> { fn main() -> anyhow::Result<()> {
@ -34,25 +35,35 @@ fn main() -> anyhow::Result<()> {
let index = Index::new(&env)?; let index = Index::new(&env)?;
let before = Instant::now();
let rtxn = env.read_txn()?; let rtxn = env.read_txn()?;
let documents_ids = index.search(&rtxn, &opt.query)?; let stdin = io::stdin();
let headers = match index.headers(&rtxn)? { let lines = match opt.query {
Some(headers) => headers, Some(query) => Box::new(once(Ok(query.to_string()))),
None => return Ok(()), None => Box::new(stdin.lock().lines()) as Box<dyn Iterator<Item = _>>,
}; };
let mut stdout = io::stdout(); for result in lines {
stdout.write_all(&headers)?; let before = Instant::now();
for id in &documents_ids { let query = result?;
if let Some(content) = index.documents.get(&rtxn, &BEU32::new(*id))? { let documents_ids = index.search(&rtxn, &query)?;
stdout.write_all(&content)?; let headers = match index.headers(&rtxn)? {
Some(headers) => headers,
None => return Ok(()),
};
let mut stdout = io::stdout();
stdout.write_all(&headers)?;
for id in &documents_ids {
if let Some(content) = index.documents.get(&rtxn, &BEU32::new(*id))? {
stdout.write_all(&content)?;
}
} }
}
eprintln!("Took {:.02?} to find {} documents", before.elapsed(), documents_ids.len()); eprintln!("Took {:.02?} to find {} documents", before.elapsed(), documents_ids.len());
}
Ok(()) Ok(())
} }