ruwex 0.1.0

//! Streaming parser for MediaWiki XML dumps.
//!
//! Works on complete dumps and on fragments (multistream blocks have no root
//! element; the final block closes a `<mediawiki>` tag that was opened in the
//! header block), so the reader is configured to be lenient about unmatched
//! end tags.

use std::io::BufRead;

use quick_xml::Reader;
use quick_xml::events::{BytesStart, BytesText, Event};

use crate::dump::{Page, SiteInfo};
use crate::error::{Error, Result};

pub struct DumpParser<R: BufRead> {
    reader: Reader<R>,
    buf: Vec<u8>,
    site_info: Option<SiteInfo>,
    pending: Option<Page>,
    finished: bool,
}

/// Top-level items the parser advances over.
enum Item {
    Page,
    SiteInfo,
    Eof,
}

/// Leaf fields captured inside `<page>`.
enum PageField {
    Title,
    Ns,
    Id,
    RevId,
    Text,
}

/// Leaf fields captured inside `<siteinfo>`.
enum SiteField {
    SiteName,
    Base,
    Namespace(i32),
}

impl<R: BufRead> DumpParser<R> {
    pub fn new(source: R) -> Self {
        let mut reader = Reader::from_reader(source);
        let config = reader.config_mut();
        config.check_end_names = false;
        config.allow_unmatched_ends = true;
        Self {
            reader,
            buf: Vec::new(),
            site_info: None,
            pending: None,
            finished: false,
        }
    }

    /// The `<siteinfo>` header, parsing forward as far as needed. `None` if
    /// the input has none (e.g. a multistream page block).
    pub fn site_info(&mut self) -> Result<Option<&SiteInfo>> {
        while self.site_info.is_none() && self.pending.is_none() && !self.finished {
            self.advance()?;
        }
        Ok(self.site_info.as_ref())
    }

    pub fn next_page(&mut self) -> Result<Option<Page>> {
        loop {
            if let Some(page) = self.pending.take() {
                return Ok(Some(page));
            }
            if self.finished {
                return Ok(None);
            }
            self.advance()?;
        }
    }

    /// Processes events until a page or the site info has been parsed, or EOF.
    fn advance(&mut self) -> Result<()> {
        let item = loop {
            self.buf.clear();
            match self.reader.read_event_into(&mut self.buf)? {
                Event::Start(e) => match e.local_name().as_ref() {
                    b"page" => break Item::Page,
                    b"siteinfo" => break Item::SiteInfo,
                    _ => {}
                },
                Event::Eof => break Item::Eof,
                _ => {}
            }
        };
        match item {
            Item::Page => self.pending = Some(self.parse_page()?),
            Item::SiteInfo => self.site_info = Some(self.parse_site_info()?),
            Item::Eof => self.finished = true,
        }
        Ok(())
    }

    /// Parses the body of a `<page>` element (the start tag is consumed).
    fn parse_page(&mut self) -> Result<Page> {
        let mut page = Page::default();
        let mut capture: Option<PageField> = None;
        let mut value = String::new();
        let mut in_revision = false;
        let mut skip_buf = Vec::new();

        loop {
            self.buf.clear();
            match self.reader.read_event_into(&mut self.buf)? {
                Event::Start(e) => {
                    let field = match e.local_name().as_ref() {
                        b"title" => Some(PageField::Title),
                        b"ns" => Some(PageField::Ns),
                        b"id" => Some(if in_revision {
                            PageField::RevId
                        } else {
                            PageField::Id
                        }),
                        b"text" => Some(PageField::Text),
                        b"revision" => {
                            in_revision = true;
                            None
                        }
                        // Skip the whole subtree so <contributor><id> is not
                        // mistaken for the revision id.
                        b"contributor" => {
                            skip_buf.clear();
                            self.reader
                                .read_to_end_into(e.to_end().name(), &mut skip_buf)?;
                            None
                        }
                        b"redirect" => {
                            page.redirect = Some(attribute(&e, "title")?);
                            None
                        }
                        _ => None,
                    };
                    if field.is_some() {
                        value.clear();
                        capture = field;
                    }
                }
                Event::Empty(e) => {
                    if e.local_name().as_ref() == b"redirect" {
                        page.redirect = Some(attribute(&e, "title")?);
                    }
                }
                Event::Text(t) => {
                    if capture.is_some() {
                        append_text(&mut value, &t);
                    }
                }
                Event::CData(c) => {
                    if capture.is_some() {
                        value.push_str(&String::from_utf8_lossy(&c));
                    }
                }
                Event::End(e) => match capture.take() {
                    Some(PageField::Title) => page.title = std::mem::take(&mut value),
                    Some(PageField::Ns) => page.ns = value.trim().parse().unwrap_or(0),
                    Some(PageField::Id) => page.id = value.trim().parse().unwrap_or(0),
                    Some(PageField::RevId) => page.revid = value.trim().parse().unwrap_or(0),
                    Some(PageField::Text) => page.text = std::mem::take(&mut value),
                    None => match e.local_name().as_ref() {
                        b"revision" => in_revision = false,
                        b"page" => return Ok(page),
                        _ => {}
                    },
                },
                Event::Eof => {
                    return Err(Error::InvalidDump("truncated <page> element".to_string()));
                }
                _ => {}
            }
        }
    }

    /// Parses the body of a `<siteinfo>` element (the start tag is consumed).
    /// Lenient about EOF so a bare dump header can be parsed on its own.
    fn parse_site_info(&mut self) -> Result<SiteInfo> {
        let mut site = SiteInfo::default();
        let mut capture: Option<SiteField> = None;
        let mut value = String::new();

        loop {
            self.buf.clear();
            match self.reader.read_event_into(&mut self.buf)? {
                Event::Start(e) => {
                    let field = match e.local_name().as_ref() {
                        b"sitename" => Some(SiteField::SiteName),
                        b"base" => Some(SiteField::Base),
                        b"namespace" => Some(SiteField::Namespace(namespace_key(&e)?)),
                        _ => None,
                    };
                    if field.is_some() {
                        value.clear();
                        capture = field;
                    }
                }
                // The main namespace has no name: <namespace key="0" />
                Event::Empty(e) => {
                    if e.local_name().as_ref() == b"namespace" {
                        site.namespaces.insert(namespace_key(&e)?, String::new());
                    }
                }
                Event::Text(t) => {
                    if capture.is_some() {
                        append_text(&mut value, &t);
                    }
                }
                Event::CData(c) => {
                    if capture.is_some() {
                        value.push_str(&String::from_utf8_lossy(&c));
                    }
                }
                Event::End(e) => match capture.take() {
                    Some(SiteField::SiteName) => site.site_name = std::mem::take(&mut value),
                    Some(SiteField::Base) => site.base = std::mem::take(&mut value),
                    Some(SiteField::Namespace(key)) => {
                        site.namespaces.insert(key, std::mem::take(&mut value));
                    }
                    None => {
                        if e.local_name().as_ref() == b"siteinfo" {
                            return Ok(site);
                        }
                    }
                },
                Event::Eof => return Ok(site),
                _ => {}
            }
        }
    }
}

impl<R: BufRead> Iterator for DumpParser<R> {
    type Item = Result<Page>;

    fn next(&mut self) -> Option<Self::Item> {
        self.next_page().transpose()
    }
}

/// Appends a text event's content, decoding XML entities. Falls back to the
/// raw bytes if the content is not strictly well-formed.
fn append_text(value: &mut String, text: &BytesText) {
    match text.unescape() {
        Ok(s) => value.push_str(&s),
        Err(_) => value.push_str(&String::from_utf8_lossy(text)),
    }
}

fn attribute(element: &BytesStart, name: &str) -> Result<String> {
    Ok(element
        .try_get_attribute(name)?
        .map(|a| a.unescape_value().map(|v| v.into_owned()))
        .transpose()?
        .unwrap_or_default())
}

fn namespace_key(element: &BytesStart) -> Result<i32> {
    Ok(attribute(element, "key")?.trim().parse().unwrap_or(0))
}

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

    const DUMP: &str = r#"<mediawiki xmlns="http://www.mediawiki.org/xml/export-0.10/" xml:lang="en">
  <siteinfo>
    <sitename>Wikipedia</sitename>
    <base>https://en.wikipedia.org/wiki/Main_Page</base>
    <namespaces>
      <namespace key="0" />
      <namespace key="1">Talk</namespace>
      <namespace key="10">Template</namespace>
    </namespaces>
  </siteinfo>
  <page>
    <title>Alpha &amp; Omega</title>
    <ns>0</ns>
    <id>1</id>
    <revision>
      <id>100</id>
      <contributor>
        <username>Someone</username>
        <id>999</id>
      </contributor>
      <text bytes="20">Hello &lt;world&gt; &amp; more</text>
    </revision>
  </page>
  <page>
    <title>Beta</title>
    <ns>0</ns>
    <id>2</id>
    <redirect title="Alpha &amp; Omega" />
    <revision>
      <id>200</id>
      <text>#REDIRECT [[Alpha &amp; Omega]]</text>
    </revision>
  </page>
</mediawiki>
"#;

    #[test]
    fn parses_site_info() {
        let mut parser = DumpParser::new(DUMP.as_bytes());
        let site = parser.site_info().unwrap().unwrap();
        assert_eq!(site.site_name, "Wikipedia");
        assert_eq!(site.base, "https://en.wikipedia.org/wiki/Main_Page");
        assert_eq!(site.namespaces[&0], "");
        assert_eq!(site.namespaces[&1], "Talk");
        assert_eq!(site.namespaces[&10], "Template");
    }

    #[test]
    fn parses_pages_after_site_info_lookup() {
        let mut parser = DumpParser::new(DUMP.as_bytes());
        parser.site_info().unwrap();

        let page = parser.next_page().unwrap().unwrap();
        assert_eq!(page.id, 1);
        assert_eq!(page.revid, 100, "contributor id must not overwrite revid");
        assert_eq!(page.ns, 0);
        assert_eq!(page.title, "Alpha & Omega");
        assert_eq!(page.text, "Hello <world> & more");
        assert_eq!(page.redirect, None);

        let page = parser.next_page().unwrap().unwrap();
        assert_eq!(page.id, 2);
        assert_eq!(page.redirect.as_deref(), Some("Alpha & Omega"));

        assert!(parser.next_page().unwrap().is_none());
    }

    #[test]
    fn parses_rootless_fragment_with_stray_end_tag() {
        let fragment = "  <page>\n    <title>Solo</title>\n    <ns>4</ns>\n    <id>7</id>\n    \
                        <revision><id>70</id><text>body</text></revision>\n  </page>\n</mediawiki>";
        let mut parser = DumpParser::new(fragment.as_bytes());
        let page = parser.next_page().unwrap().unwrap();
        assert_eq!(page.id, 7);
        assert_eq!(page.ns, 4);
        assert_eq!(page.text, "body");
        assert!(parser.next_page().unwrap().is_none());
        assert!(parser.site_info().unwrap().is_none());
    }

    #[test]
    fn parses_bare_header_without_closing_root() {
        let header = "<mediawiki version=\"0.10\">\n  <siteinfo>\n    <sitename>Wiki</sitename>\n    \
                      <base>http://example.org/wiki/Main</base>\n    <namespaces><namespace key=\"0\" /></namespaces>\n  </siteinfo>\n";
        let mut parser = DumpParser::new(header.as_bytes());
        let site = parser.site_info().unwrap().unwrap();
        assert_eq!(site.site_name, "Wiki");
        assert_eq!(site.url_base(), "http://example.org/wiki");
    }
}