ruwex 0.1.0

Fast Rust rewrite of wikiextractor: extract and clean text from Wikimedia XML dumps
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125

//! Port of wikiextractor's `cirrus-extract.py`: extracts documents from
//! Wikipedia CirrusSearch JSON dumps, whose text already has templates
//! expanded. Output matches the original byte for byte — including its
//! hardcoded `http://it.wikipedia.org/` URL base, kept for parity.

use std::io::BufRead;
use std::sync::LazyLock;

use regex::Regex;

use crate::error::{Error, Result};
use crate::output::writer::DocSink;

/// Reference lines like `  ^ The Penguin Dictionary` are dropped.
static REFERENCE_LINE: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"  \^ .*").unwrap());

/// Processes a CirrusSearch dump: alternating index/content JSON lines.
/// Returns the number of documents written.
pub fn process_cirrus_dump(input: impl BufRead, sink: &mut dyn DocSink) -> Result<u64> {
    let mut docs = 0;
    let mut lines = input.lines();
    while let Some(index_line) = lines.next() {
        let index_line = index_line?;
        if index_line.trim().is_empty() {
            continue;
        }
        let index: serde_json::Value = serde_json::from_str(&index_line)
            .map_err(|e| Error::InvalidDump(format!("bad cirrus index line: {e}")))?;
        let Some(content_line) = lines.next() else {
            return Err(Error::InvalidDump(
                "cirrus dump ends after an index line".to_string(),
            ));
        };
        let content: serde_json::Value = serde_json::from_str(&content_line?)
            .map_err(|e| Error::InvalidDump(format!("bad cirrus content line: {e}")))?;

        // upstream checks _type == "page"; newer Cirrus dumps use "_doc"
        let doc_type = index["index"]["_type"].as_str().unwrap_or("");
        if !matches!(doc_type, "page" | "_doc") || content["namespace"] != 0 {
            continue;
        }
        let id = json_scalar(&index["index"]["_id"]);
        let language = json_scalar(&content["language"]);
        let revision = json_scalar(&content["version"]);
        let title = content["title"].as_str().unwrap_or("");
        let text = content["text"].as_str().unwrap_or("");
        let text = REFERENCE_LINE.replace_all(text, "");

        let doc = format!(
            "<doc id=\"{id}\" url=\"http://it.wikipedia.org/wiki?curid={id}\" \
             title=\"{title}\" language=\"{language}\" revision=\"{revision}\">\n\
             {title}\n\n{text}\n</doc>\n"
        );
        sink.write_doc(&doc)?;
        docs += 1;
    }
    Ok(docs)
}

/// Renders a JSON scalar the way Python's `%s` would (strings unquoted,
/// numbers as digits, missing values empty).
fn json_scalar(value: &serde_json::Value) -> String {
    match value {
        serde_json::Value::String(s) => s.clone(),
        serde_json::Value::Null => String::new(),
        other => other.to_string(),
    }
}

#[cfg(test)]
mod tests {
    use std::io;

    use super::*;

    #[derive(Default)]
    struct StringSink(String);

    impl DocSink for StringSink {
        fn write_doc(&mut self, doc: &str) -> io::Result<()> {
            self.0.push_str(doc);
            Ok(())
        }

        fn finish(&mut self) -> io::Result<()> {
            Ok(())
        }
    }

    fn dump() -> String {
        [
            r#"{"index":{"_type":"page","_id":"3825914"}}"#,
            r#"{"namespace":0,"title":"Alpha","language":"en","version":42,"text":"Alpha body.  ^ A reference\nMore."}"#,
            r#"{"index":{"_type":"page","_id":"77"}}"#,
            r#"{"namespace":4,"title":"Project:Skip","language":"en","version":1,"text":"nope"}"#,
            r#"{"index":{"_type":"_doc","_id":"99"}}"#,
            r#"{"namespace":0,"title":"Beta","language":"en","version":7,"text":"Beta body."}"#,
        ]
        .join("\n")
    }

    #[test]
    fn extracts_main_namespace_docs() {
        let mut sink = StringSink::default();
        let docs = process_cirrus_dump(dump().as_bytes(), &mut sink).unwrap();
        assert_eq!(docs, 2);
        assert_eq!(
            sink.0,
            "<doc id=\"3825914\" url=\"http://it.wikipedia.org/wiki?curid=3825914\" \
             title=\"Alpha\" language=\"en\" revision=\"42\">\nAlpha\n\nAlpha body.\nMore.\n</doc>\n\
             <doc id=\"99\" url=\"http://it.wikipedia.org/wiki?curid=99\" \
             title=\"Beta\" language=\"en\" revision=\"7\">\nBeta\n\nBeta body.\n</doc>\n"
        );
    }

    #[test]
    fn truncated_dump_is_an_error() {
        let mut sink = StringSink::default();
        let result = process_cirrus_dump(
            r#"{"index":{"_type":"page","_id":"1"}}"#.as_bytes(),
            &mut sink,
        );
        assert!(result.is_err());
    }
}