ruwex 0.1.0

Fast Rust rewrite of wikiextractor: extract and clean text from Wikimedia XML dumps
Documentation
use crate::error::{Error, Result};

/// Format of the extracted documents.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum OutputFormat {
    /// `<doc id=... url=... title=...>` blocks (wikiextractor default).
    #[default]
    Doc,
    /// One JSON object per line, in wikiextractor's exact shape.
    Json,
}

/// Extraction settings shared by the library pipeline and the CLI.
#[derive(Debug, Clone)]
pub struct ExtractorConfig {
    pub format: OutputFormat,
    /// Preserve links as `<a>` elements in the cleaned text.
    pub keep_links: bool,
    /// Keep basic HTML formatting (`<b>`, `<i>`, headings, lists) in the
    /// output; implies `keep_links` in the CLI.
    pub html: bool,
    /// Escape `&`, `<`, `>` in the extracted text so it is safe inside
    /// `<doc>…</doc>`.
    pub html_safe: bool,
    /// Accepted namespace prefixes. Pages whose title starts with `ns:` are
    /// kept only if `ns` is listed here, and internal links into other
    /// namespaces are dropped. The default matches wikiextractor:
    /// interlanguage-style prefixes only, so effectively main-namespace
    /// articles.
    pub namespaces: Vec<String>,
    /// Number of extraction worker threads.
    pub workers: usize,
}

impl Default for ExtractorConfig {
    fn default() -> Self {
        Self {
            format: OutputFormat::default(),
            keep_links: false,
            html: false,
            html_safe: true,
            namespaces: default_namespaces(),
            workers: default_workers(),
        }
    }
}

/// wikiextractor's default `acceptedNamespaces`: cross-wiki link prefixes.
pub fn default_namespaces() -> Vec<String> {
    ["w", "wiktionary", "wikt"].map(String::from).to_vec()
}

/// Default worker count, matching wikiextractor: all cores but one.
pub fn default_workers() -> usize {
    std::thread::available_parallelism()
        .map(|n| n.get())
        .unwrap_or(2)
        .saturating_sub(1)
        .max(1)
}

/// Parses a size argument in wikiextractor's `n[KMG]` notation into bytes.
/// `0` is accepted and means "no limit" to callers.
pub fn parse_size(size: &str) -> Result<u64> {
    let size = size.trim();
    let (digits, multiplier) = match size.as_bytes().last() {
        Some(b'k' | b'K') => (&size[..size.len() - 1], 1u64 << 10),
        Some(b'm' | b'M') => (&size[..size.len() - 1], 1u64 << 20),
        Some(b'g' | b'G') => (&size[..size.len() - 1], 1u64 << 30),
        _ => (size, 1),
    };
    digits
        .parse::<u64>()
        .ok()
        .and_then(|n| n.checked_mul(multiplier))
        .ok_or_else(|| Error::InvalidSize(size.to_string()))
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn parses_plain_and_suffixed_sizes() {
        assert_eq!(parse_size("123").unwrap(), 123);
        assert_eq!(parse_size("500K").unwrap(), 500 * 1024);
        assert_eq!(parse_size("1M").unwrap(), 1024 * 1024);
        assert_eq!(parse_size("2g").unwrap(), 2 * 1024 * 1024 * 1024);
        assert_eq!(parse_size(" 1M ").unwrap(), 1024 * 1024);
        assert_eq!(parse_size("0").unwrap(), 0);
    }

    #[test]
    fn rejects_invalid_sizes() {
        assert!(parse_size("").is_err());
        assert!(parse_size("M").is_err());
        assert!(parse_size("12X").is_err());
        assert!(parse_size("-5K").is_err());
    }
}