ruwex 0.1.0

Fast Rust rewrite of wikiextractor: extract and clean text from Wikimedia XML dumps
Documentation
//! Template expansion, ported from wikiextractor's `Extractor.expandTemplates`:
//! magic words, parser functions, and template instantiation against a
//! [`TemplateDb`]. With an empty database (`--no-templates`) every ordinary
//! inclusion resolves to the empty string, exactly like the original.

pub mod braces;
pub mod functions;
pub mod template;
pub mod templates;

use std::collections::HashMap;
use std::sync::Arc;

use braces::find_matching_braces;
use functions::call_parser_function;
use std::sync::LazyLock;

use regex::Regex;
use template::Template;
pub use templates::{LazyTemplateSource, TemplateDb};

/// Python's `maxTemplateRecursionLevels`, checked against the frame depth.
const MAX_TEMPLATE_RECURSION: usize = 30;

/// A source of template definitions the [`Expander`] resolves invocations
/// against. Implemented both by [`TemplateDb`] (all definitions loaded up
/// front — right for bulk extraction) and by [`LazyTemplateSource`] (fetched
/// on demand through a title index — right for a single-page lookup).
pub trait TemplateSource {
    /// Resolves a raw invocation title to a fully-qualified page title.
    fn qualify(&self, title: &str) -> String;
    /// Returns the parsed template for a qualified title, or `None` if there
    /// is no such template (following redirects).
    fn get_parsed(&self, qualified_title: &str) -> Option<Arc<Template>>;
}

pub struct Expander<'a> {
    /// Lowercase magic-word name → value (per page, like Python's MagicWords).
    magic_words: HashMap<&'static str, String>,
    /// Invocation frames `(qualified title, parameters)`; used both for the
    /// recursion limit and for zero-progress loop detection.
    frame: Vec<(String, HashMap<String, String>)>,
    source: &'a dyn TemplateSource,
}

impl<'a> Expander<'a> {
    pub fn new(page_title: &str, source: &'a dyn TemplateSource) -> Self {
        let now = jiff::Zoned::now();
        let strftime = |fmt: &str| now.strftime(fmt).to_string();
        let namespace = match page_title.find(':') {
            Some(colon) => page_title[..colon].to_string(),
            None => String::new(),
        };
        let magic_words = HashMap::from([
            ("!", "|".to_string()),
            ("namespace", namespace),
            ("pagename", page_title.to_string()),
            ("fullpagename", page_title.to_string()),
            ("currentyear", strftime("%Y")),
            ("currentmonth", strftime("%m")),
            ("currentday", strftime("%d")),
            ("currenthour", strftime("%H")),
            ("currenttime", strftime("%H:%M:%S")),
        ]);
        Self {
            magic_words,
            frame: Vec::new(),
            source,
        }
    }

    /// Expands all `{{…}}` / `{{{…}}}` constructs in `text`.
    pub fn expand(&mut self, text: &str) -> String {
        if self.frame.len() >= MAX_TEMPLATE_RECURSION {
            log::debug!("max template recursion exceeded");
            return String::new();
        }
        let mut result = String::with_capacity(text.len());
        let mut cur = 0;
        for (s, e) in find_matching_braces(text, 2) {
            result.push_str(&text[cur..s]);
            // defensive lenient slice, mirroring Python (see Template::parse)
            let body = text.get(s + 2..e.saturating_sub(2)).unwrap_or("");
            result.push_str(&self.expand_invocation(body));
            cur = e;
        }
        result.push_str(&text[cur..]);
        result
    }

    /// Expands one invocation (the text between the braces).
    fn expand_invocation(&mut self, body: &str) -> String {
        if self.frame.len() >= MAX_TEMPLATE_RECURSION {
            log::debug!("max template recursion exceeded");
            return String::new();
        }
        let parts = split_parts(body);
        let expanded_title = self.expand(parts[0].trim());
        let (title, subst) = strip_subst(&expanded_title);

        if let Some(value) = self.magic_words.get(title.to_lowercase().as_str()) {
            return value.clone();
        }

        // Parser functions: the first argument is everything after the colon.
        if let Some(colon) = title.find(':')
            && title[..colon].chars().count() > 1
        {
            let function = &title[..colon];
            let mut args: Vec<String> = parts.iter().map(|p| p.to_string()).collect();
            args[0] = title[colon + 1..].trim().to_string();
            let result = call_parser_function(function, &args);
            return self.expand(&result);
        }

        let qualified = self.source.qualify(title);
        if qualified.is_empty() {
            log::debug!("empty template title in invocation");
            return String::new();
        }
        let Some(template) = self.source.get_parsed(&qualified) else {
            // page being included could not be identified
            return String::new();
        };

        // Evaluate parameters (unless subst), since they may contain
        // templates — including the symbol "=".
        let mut params: Vec<String> = parts[1..].to_vec();
        if !subst {
            params = params.iter().map(|p| self.expand(p)).collect();
        }
        let params = template_params(&params);

        // Zero-progress loop detection: the same title invoked again with
        // the same parameters as an active ancestor is left unexpanded.
        if self
            .frame
            .iter()
            .any(|(frame_title, frame_params)| *frame_title == qualified && *frame_params == params)
        {
            log::debug!("template loop detected: {qualified}");
            return format!("{{{{{body}}}}}");
        }

        self.frame.push((qualified, params.clone()));
        let instantiated = template.subst(&params, self, 0);
        let value = self.expand(&instantiated);
        self.frame.pop();
        value
    }
}

/// Removes a leading `subst:` / `safesubst:` (case-insensitive); the flag
/// reports whether one was present.
fn strip_subst(title: &str) -> (&str, bool) {
    for prefix in ["subst:", "safesubst:"] {
        // compare as bytes: a byte-offset slice of the str could split a
        // multibyte character; the prefix itself is pure ASCII, so a
        // case-insensitive byte match can only succeed on an ASCII prefix
        if title.len() >= prefix.len()
            && title.as_bytes()[..prefix.len()].eq_ignore_ascii_case(prefix.as_bytes())
        {
            return (&title[prefix.len()..], true);
        }
    }
    (title, false)
}

static NAMED_PARAM: LazyLock<Regex> =
    LazyLock::new(|| Regex::new(r"(?s)^ *([^=']*?) *=(.*)").unwrap());

/// Port of `templateParams`: builds the name/position → value map. Values
/// containing a link (`]]`) keep their surrounding whitespace.
fn template_params(params: &[String]) -> HashMap<String, String> {
    let mut map = HashMap::new();
    let mut unnamed = 0;
    for param in params {
        if let Some(caps) = NAMED_PARAM.captures(param) {
            let name = caps[1].trim().to_string();
            let mut value = caps.get(2).expect("value group").as_str();
            if !value.contains("]]") {
                value = value.trim();
            }
            map.insert(name, value.to_string());
        } else {
            unnamed += 1;
            let value = if param.contains("]]") {
                param.as_str()
            } else {
                param.trim()
            };
            map.insert(unnamed.to_string(), value.to_string());
        }
    }
    map
}

/// Splits template parameters at `|` symbols that are not inside any
/// `{{…}}`, `{{{…}}}` or `[[…]]` pair (port of `splitParts`).
pub fn split_parts(params_list: &str) -> Vec<String> {
    let mut parameters: Vec<String> = Vec::new();
    let mut cur = 0;
    let append_split = |parameters: &mut Vec<String>, segment: &str| {
        let mut pieces = segment.split('|');
        let first = pieces.next().expect("split yields at least one piece");
        match parameters.last_mut() {
            // portion before | belongs to previous parameter
            Some(last) => last.push_str(first),
            None => parameters.push(first.to_string()),
        }
        parameters.extend(pieces.map(str::to_string));
    };
    for (s, e) in find_matching_braces(params_list, 0) {
        append_split(&mut parameters, &params_list[cur..s]);
        // the balanced span belongs to the last parameter as-is
        parameters
            .last_mut()
            .expect("at least one parameter exists")
            .push_str(&params_list[s..e]);
        cur = e;
    }
    append_split(&mut parameters, &params_list[cur..]);
    parameters
}

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

    fn expand(text: &str) -> String {
        let db = TemplateDb::default();
        Expander::new("Test page", &db).expand(text)
    }

    fn expand_with(templates: &[(&str, &str)], text: &str) -> String {
        let mut db = TemplateDb::default();
        for (title, body) in templates {
            db.define(title, body);
        }
        Expander::new("Test page", &db).expand(text)
    }

    #[test]
    fn unknown_templates_expand_to_nothing() {
        assert_eq!(expand("a {{lang-grc|x|y}} b"), "a  b");
        assert_eq!(expand("a {{nested|x={{deep|y}}}} b"), "a  b");
    }

    #[test]
    fn top_level_template_arguments_vanish() {
        assert_eq!(expand("a {{{1|default}}} b"), "a  b");
    }

    #[test]
    fn magic_words_expand() {
        assert_eq!(expand("{{PAGENAME}}"), "Test page");
        assert_eq!(expand("a{{!}}b"), "a|b");
        let db = TemplateDb::default();
        assert_eq!(Expander::new("Talk:X", &db).expand("{{NAMESPACE}}"), "Talk");
    }

    #[test]
    fn parser_functions_evaluate() {
        assert_eq!(expand("{{#if:x|yes|no}}"), "yes");
        assert_eq!(expand("{{#if:|yes|no}}"), "no");
        assert_eq!(expand("{{#expr:2*3+4}}"), "10");
        assert_eq!(expand("{{lc:ABC}}"), "abc");
        assert_eq!(expand("{{#switch:b|a=1|b=2}}"), "2");
    }

    #[test]
    fn branches_are_expanded_lazily_from_results() {
        // the winning branch is expanded after selection
        assert_eq!(expand("{{#if:x|{{#ifeq:a|a|inner}}|other}}"), "inner");
        assert_eq!(expand("{{#if:x|{{unknown template}}rest}}"), "rest");
    }

    #[test]
    fn subst_prefix_is_stripped() {
        assert_eq!(expand("{{subst:lc:ABC}}"), "abc");
        assert_eq!(expand("{{safesubst:PAGENAME}}"), "Test page");
        // multibyte titles must not panic the ASCII prefix check
        assert_eq!(expand("{{lang\u{2013}x|y}}"), "");
        assert_eq!(expand("{{\u{fc}ber}}"), "");
    }

    #[test]
    fn split_parts_respects_nesting() {
        assert_eq!(
            split_parts("t|a|b={{x|y}}|c [[l|lab]] d"),
            vec!["t", "a", "b={{x|y}}", "c [[l|lab]] d"]
        );
        assert_eq!(split_parts(""), vec![""]);
        assert_eq!(split_parts("a||c"), vec!["a", "", "c"]);
    }

    #[test]
    fn defined_templates_instantiate() {
        let templates = [("Template:Greet", "Hello {{{name|stranger}}}!")];
        assert_eq!(
            expand_with(&templates, "{{Greet|name=World}}"),
            "Hello World!"
        );
        assert_eq!(expand_with(&templates, "{{greet}}"), "Hello stranger!");
    }

    #[test]
    fn nested_template_invocations_expand() {
        let templates = [
            ("Template:Inner", "[{{{1}}}]"),
            ("Template:Outer", "out {{Inner|{{{x}}}}} out"),
        ];
        assert_eq!(expand_with(&templates, "{{Outer|x=42}}"), "out [42] out");
    }

    #[test]
    fn parameters_may_contain_templates() {
        let templates = [("Template:Id", "{{{1}}}")];
        assert_eq!(expand_with(&templates, "{{Id|{{lc:ABC}}}}"), "abc");
    }

    #[test]
    fn later_named_parameter_overrides_positional() {
        let templates = [("Template:T", "{{{1}}}-{{{2}}}-{{{3}}}")];
        // {{t|a|b|c|2=B}} is equivalent to {{t|a|B|c}}
        assert_eq!(expand_with(&templates, "{{T|a|b|c|2=B}}"), "a-B-c");
    }

    #[test]
    fn template_redirects_resolve() {
        let templates = [
            ("Template:Real", "value"),
            ("Template:Alias", "#REDIRECT [[Template:Real]]"),
        ];
        assert_eq!(expand_with(&templates, "{{Alias}}"), "value");
    }

    #[test]
    fn zero_progress_loops_stay_literal() {
        let templates = [("Template:Loop", "x{{Loop}}")];
        assert_eq!(expand_with(&templates, "{{Loop}}"), "x{{Loop}}");
    }

    #[test]
    fn self_recursion_with_different_params_progresses() {
        // a template iterating by shrinking its argument list must not be
        // blocked by loop detection
        let templates = [(
            "Template:Count",
            "{{{1|}}}{{#if:{{{2|}}}|{{Count|{{{2}}}|{{{3|}}}}}}}",
        )];
        assert_eq!(expand_with(&templates, "{{Count|a|b|c}}"), "abc");
    }

    #[test]
    fn recursion_limit_returns_empty() {
        // infinite self-inclusion with changing params exhausts the frame limit
        let templates = [(
            "Template:Inf",
            "{{#expr:{{{1}}}+1}}{{Inf|{{#expr:{{{1}}}+1}}}}",
        )];
        let result = expand_with(&templates, "{{Inf|0}}");
        // must terminate; the tail is cut off at the recursion limit
        assert!(result.starts_with("123"), "unexpected: {result}");
    }
}