ruwex 0.1.0

//! Drop-in replacement for wikiextractor's `cirrus-extract.py`: extracts
//! documents from Wikipedia CirrusSearch JSON dumps.

use std::ffi::OsString;
use std::fs::File;
use std::io::{self, BufRead, BufReader};
use std::path::Path;

use anyhow::{Context, ensure};
use bzip2::read::MultiBzDecoder;
use clap::Parser;
use flate2::bufread::MultiGzDecoder;
use log::{LevelFilter, info};

use ruwex::tools::cirrus::process_cirrus_dump;
use ruwex::{DocSink, ShardedWriter, StdoutSink, parse_size};

/// Minimum size of output files, as in the original.
const MIN_FILE_SIZE: u64 = 200 * 1024;

#[derive(Parser, Debug)]
#[command(
    name = "cirrus-extract",
    about = "Extracts and cleans text from a Wikipedia Cirrus dump. \
             Rust rewrite (ruwex) of cirrus-extract.py, CLI-compatible.",
    disable_version_flag = true
)]
struct Args {
    /// Cirrus Json wiki dump file
    input: Option<String>,

    /// directory for extracted files (or '-' for dumping to stdout)
    #[arg(short, long, default_value = "text")]
    output: String,

    /// maximum bytes per output file
    #[arg(short, long, default_value = "1M", value_name = "n[KMG]")]
    bytes: String,

    /// compress output files using bzip
    #[arg(short, long)]
    compress: bool,

    /// accepted namespaces (accepted for compatibility; the original
    /// ignores it and keeps the main namespace only)
    #[arg(long, value_name = "ns1,ns2")]
    namespaces: Option<String>,

    /// suppress reporting progress info
    #[arg(short, long)]
    quiet: bool,

    /// print program version
    #[arg(short = 'v', long = "version")]
    version: bool,
}

/// argparse accepts `-ns`; map it to the long form before parsing.
fn preprocessed_args() -> Vec<OsString> {
    std::env::args_os()
        .map(|arg| {
            if arg == "-ns" {
                "--namespaces".into()
            } else {
                arg
            }
        })
        .collect()
}

fn main() {
    if let Err(error) = run() {
        eprintln!("error: {error:#}");
        std::process::exit(1);
    }
}

fn run() -> anyhow::Result<()> {
    let args = Args::parse_from(preprocessed_args());
    if args.version {
        println!("cirrus-extract {} (ruwex)", env!("CARGO_PKG_VERSION"));
        return Ok(());
    }
    let input = args
        .input
        .context("the following arguments are required: input")?;

    env_logger::Builder::new()
        .filter_level(if args.quiet {
            LevelFilter::Warn
        } else {
            LevelFilter::Info
        })
        .format_target(false)
        .init();

    let max_bytes = parse_size(&args.bytes)?;
    ensure!(
        max_bytes >= MIN_FILE_SIZE,
        "Insufficient or invalid size: {}",
        args.bytes
    );

    let reader: Box<dyn BufRead> = if input == "-" {
        Box::new(BufReader::new(io::stdin()))
    } else {
        open_compressed(&input).with_context(|| format!("cannot open {input}"))?
    };

    let mut sink: Box<dyn DocSink> = if args.output == "-" {
        Box::new(StdoutSink::new())
    } else {
        Box::new(
            ShardedWriter::create(Path::new(&args.output), max_bytes, args.compress)
                .with_context(|| format!("cannot create output in {}", args.output))?,
        )
    };

    let docs = process_cirrus_dump(reader, sink.as_mut())?;
    sink.finish()?;
    info!("finished: {docs} documents written");
    Ok(())
}

/// Opens a cirrus dump, sniffing gzip (the usual packaging) or bz2 by magic
/// bytes; anything else is read as plain text.
fn open_compressed(path: &str) -> io::Result<Box<dyn BufRead>> {
    let mut reader = BufReader::new(File::open(path)?);
    let magic = reader.fill_buf()?;
    Ok(if magic.starts_with(&[0x1f, 0x8b]) {
        Box::new(BufReader::new(MultiGzDecoder::new(reader)))
    } else if magic.starts_with(b"BZh") {
        Box::new(BufReader::new(MultiBzDecoder::new(reader)))
    } else {
        Box::new(reader)
    })
}