Benjamin Mako Hill
3d6bf34c8e
A single mwparserfromhell tree walk per revision emits two nullable int64 columns. prose_chars counts rendered content outside any template or tag: text nodes, wikilink labels (or targets when unlabeled), external link labels, and heading titles. structured_chars counts the same kinds of content inside template parameter values or tag contents, with each character assigned by its nearest enclosing container. Markup syntax, template and tag names, parameter names, attributes, bare URLs, and HTML comments count in neither, so syntax overhead is derivable as the total size minus the two columns. Deleted revisions and parse timeouts yield nulls through the same machinery as the other parser-based columns and do not change which revisions are processed. Note that mwparserfromhell parses bold and italic markup as tags, so emphasized text counts as structured. Includes hand-computed fixture tests for the counting rules and an end-to-end test on a dump with deleted revisions verifying nulls and the invariant prose_chars + structured_chars <= revision size. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
864 lines
34 KiB
Python
864 lines
34 KiB
Python
import os
|
|
import subprocess
|
|
import sys
|
|
import tracemalloc
|
|
from io import StringIO
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
import pyarrow as pa
|
|
import pyarrow.json as pj
|
|
import pytest
|
|
from pandas import DataFrame
|
|
from pandas.testing import assert_frame_equal, assert_series_equal
|
|
|
|
from wikiq import build_table, build_schema, RegexPair, WikitextParser
|
|
from wikiq_test_utils import (
|
|
BASELINE_DIR,
|
|
IKWIKI,
|
|
REGEXTEST,
|
|
SAILORMOON,
|
|
TEST_DIR,
|
|
TEST_OUTPUT_DIR,
|
|
TWINPEAKS,
|
|
WIKIQ,
|
|
WikiqTester,
|
|
)
|
|
|
|
|
|
def setup():
|
|
tracemalloc.start()
|
|
|
|
if not os.path.exists(TEST_OUTPUT_DIR):
|
|
os.mkdir(TEST_OUTPUT_DIR)
|
|
|
|
|
|
setup()
|
|
|
|
|
|
def read_jsonl_with_schema(filepath: str, **schema_kwargs) -> pd.DataFrame:
|
|
"""Read JSONL file using PyArrow with explicit schema from wikiq."""
|
|
table, _, _ = build_table(**schema_kwargs)
|
|
schema = build_schema(table, **schema_kwargs)
|
|
pa_table = pj.read_json(
|
|
filepath,
|
|
parse_options=pj.ParseOptions(explicit_schema=schema),
|
|
)
|
|
return pa_table.to_pandas()
|
|
|
|
|
|
# with / without pwr DONE
|
|
# with / without url encode DONE
|
|
# with / without collapse user DONE
|
|
# with output to stdout DONE
|
|
# note that the persistence radius is 7 by default
|
|
# reading various file formats including
|
|
# 7z, gz, bz2, xml DONE
|
|
# wikia and wikipedia data DONE
|
|
# malformed xmls DONE
|
|
|
|
def test_WP_noargs():
|
|
tester = WikiqTester(IKWIKI, "noargs")
|
|
|
|
try:
|
|
tester.call_wikiq()
|
|
except subprocess.CalledProcessError as exc:
|
|
pytest.fail(exc.stderr.decode("utf8"))
|
|
|
|
test = pd.read_table(tester.output)
|
|
baseline = pd.read_table(tester.baseline_file)
|
|
assert_frame_equal(test, baseline, check_like=True)
|
|
|
|
def test_WP_namespaces():
|
|
tester = WikiqTester(IKWIKI, "namespaces")
|
|
|
|
try:
|
|
tester.call_wikiq("-n 0", "-n 1")
|
|
except subprocess.CalledProcessError as exc:
|
|
pytest.fail(exc.stderr.decode("utf8"))
|
|
|
|
# as a test let's make sure that we get equal data frames
|
|
test = pd.read_table(tester.output)
|
|
num_wrong_ns = sum(~test.namespace.isin({0, 1}))
|
|
assert num_wrong_ns == 0
|
|
baseline = pd.read_table(tester.baseline_file)
|
|
assert_frame_equal(test, baseline, check_like=True)
|
|
|
|
def test_WP_revert_radius():
|
|
tester = WikiqTester(IKWIKI, "revert_radius")
|
|
|
|
try:
|
|
tester.call_wikiq("-n 0", "-n 1", "-rr 1")
|
|
except subprocess.CalledProcessError as exc:
|
|
pytest.fail(exc.stderr.decode("utf8"))
|
|
|
|
# as a test let's make sure that we get equal data frames
|
|
test = pd.read_table(tester.output)
|
|
num_wrong_ns = sum(~test.namespace.isin({0, 1}))
|
|
assert num_wrong_ns == 0
|
|
baseline = pd.read_table(tester.baseline_file)
|
|
assert_frame_equal(test, baseline, check_like=True)
|
|
|
|
def test_WP_no_revert_radius():
|
|
tester = WikiqTester(IKWIKI, "no_revert_radius")
|
|
|
|
try:
|
|
tester.call_wikiq("-rr 0")
|
|
except subprocess.CalledProcessError as exc:
|
|
pytest.fail(exc.stderr.decode("utf8"))
|
|
|
|
# as a test let's make sure that we get equal data frames
|
|
test = pd.read_table(tester.output)
|
|
num_reverted = sum(i is None for i in test.revert)
|
|
assert num_reverted == 0
|
|
baseline = pd.read_table(tester.baseline_file)
|
|
assert_frame_equal(test, baseline, check_like=True)
|
|
|
|
def test_WP_collapse_user():
|
|
tester = WikiqTester(IKWIKI, "collapse_user")
|
|
|
|
try:
|
|
tester.call_wikiq("--collapse-user")
|
|
except subprocess.CalledProcessError as exc:
|
|
pytest.fail(exc.stderr.decode("utf8"))
|
|
|
|
test = pd.read_table(tester.output)
|
|
baseline = pd.read_table(tester.baseline_file)
|
|
assert_frame_equal(test, baseline, check_like=True)
|
|
|
|
def test_noargs():
|
|
tester = WikiqTester(SAILORMOON, "noargs", in_compression="7z")
|
|
|
|
try:
|
|
tester.call_wikiq()
|
|
except subprocess.CalledProcessError as exc:
|
|
pytest.fail(exc.stderr.decode("utf8"))
|
|
|
|
test = pd.read_table(tester.output)
|
|
baseline = pd.read_table(tester.baseline_file)
|
|
assert_frame_equal(test, baseline, check_like=True)
|
|
|
|
|
|
def test_jsonl_noargs():
|
|
"""Test JSONL output format with baseline comparison."""
|
|
tester = WikiqTester(SAILORMOON, "noargs", in_compression="7z", out_format="jsonl", baseline_format="jsonl")
|
|
|
|
try:
|
|
tester.call_wikiq()
|
|
except subprocess.CalledProcessError as exc:
|
|
pytest.fail(exc.stderr.decode("utf8"))
|
|
|
|
test = read_jsonl_with_schema(tester.output)
|
|
baseline = read_jsonl_with_schema(tester.baseline_file)
|
|
assert_frame_equal(test, baseline, check_like=True)
|
|
|
|
|
|
def test_jsonl_tsv_equivalence():
|
|
"""Test that JSONL and TSV outputs contain equivalent data."""
|
|
tester_tsv = WikiqTester(SAILORMOON, "equiv_tsv", in_compression="7z", out_format="tsv")
|
|
tester_jsonl = WikiqTester(SAILORMOON, "equiv_jsonl", in_compression="7z", out_format="jsonl")
|
|
|
|
try:
|
|
tester_tsv.call_wikiq()
|
|
tester_jsonl.call_wikiq()
|
|
except subprocess.CalledProcessError as exc:
|
|
pytest.fail(exc.stderr.decode("utf8"))
|
|
|
|
tsv_df = pd.read_table(tester_tsv.output)
|
|
jsonl_df = read_jsonl_with_schema(tester_jsonl.output)
|
|
|
|
# Row counts must match
|
|
assert len(tsv_df) == len(jsonl_df), f"Row count mismatch: TSV={len(tsv_df)}, JSONL={len(jsonl_df)}"
|
|
|
|
# Column sets must match
|
|
assert set(tsv_df.columns) == set(jsonl_df.columns), \
|
|
f"Column mismatch: TSV={set(tsv_df.columns)}, JSONL={set(jsonl_df.columns)}"
|
|
|
|
# Sort both by revid for comparison
|
|
tsv_df = tsv_df.sort_values("revid").reset_index(drop=True)
|
|
jsonl_df = jsonl_df.sort_values("revid").reset_index(drop=True)
|
|
|
|
# Normalize null values: TSV uses nan, schema-based JSONL uses None
|
|
jsonl_df = jsonl_df.replace({None: np.nan})
|
|
|
|
# Compare columns - schema-based reading handles types correctly
|
|
for col in tsv_df.columns:
|
|
if col == "date_time":
|
|
# TSV reads as string, JSONL with schema reads as datetime
|
|
tsv_dates = pd.to_datetime(tsv_df[col]).dt.strftime("%Y-%m-%dT%H:%M:%SZ")
|
|
jsonl_dates = jsonl_df[col].dt.strftime("%Y-%m-%dT%H:%M:%SZ")
|
|
assert_series_equal(tsv_dates, jsonl_dates, check_names=False)
|
|
else:
|
|
# Allow dtype differences (TSV infers int64, schema uses int32)
|
|
assert_series_equal(tsv_df[col], jsonl_df[col], check_names=False, check_dtype=False)
|
|
|
|
|
|
def test_collapse_user():
|
|
tester = WikiqTester(SAILORMOON, "collapse-user", in_compression="7z")
|
|
|
|
try:
|
|
tester.call_wikiq("--collapse-user", "--fandom-2020")
|
|
except subprocess.CalledProcessError as exc:
|
|
pytest.fail(exc.stderr.decode("utf8"))
|
|
|
|
test = pd.read_table(tester.output)
|
|
baseline = pd.read_table(tester.baseline_file)
|
|
assert_frame_equal(test, baseline, check_like=True)
|
|
|
|
def test_pwr_wikidiff2():
|
|
tester = WikiqTester(SAILORMOON, "persistence_wikidiff2", in_compression="7z")
|
|
|
|
try:
|
|
tester.call_wikiq("--persistence wikidiff2", "--fandom-2020")
|
|
except subprocess.CalledProcessError as exc:
|
|
pytest.fail(exc.stderr.decode("utf8"))
|
|
|
|
test = pd.read_table(tester.output)
|
|
baseline = pd.read_table(tester.baseline_file)
|
|
assert_frame_equal(test, baseline, check_like=True)
|
|
|
|
def test_pwr_segment():
|
|
tester = WikiqTester(SAILORMOON, "persistence_segment", in_compression="7z")
|
|
|
|
try:
|
|
tester.call_wikiq("--persistence segment", "--fandom-2020")
|
|
except subprocess.CalledProcessError as exc:
|
|
pytest.fail(exc.stderr.decode("utf8"))
|
|
|
|
test = pd.read_table(tester.output)
|
|
baseline = pd.read_table(tester.baseline_file)
|
|
assert_frame_equal(test, baseline, check_like=True)
|
|
|
|
def test_pwr_legacy():
|
|
tester = WikiqTester(SAILORMOON, "persistence_legacy", in_compression="7z")
|
|
|
|
try:
|
|
tester.call_wikiq("--persistence legacy", "--fandom-2020")
|
|
except subprocess.CalledProcessError as exc:
|
|
pytest.fail(exc.stderr.decode("utf8"))
|
|
|
|
test = pd.read_table(tester.output)
|
|
baseline = pd.read_table(tester.baseline_file)
|
|
assert_frame_equal(test, baseline, check_like=True)
|
|
|
|
def test_pwr():
|
|
tester = WikiqTester(SAILORMOON, "persistence", in_compression="7z")
|
|
|
|
try:
|
|
tester.call_wikiq("--persistence", "--fandom-2020")
|
|
except subprocess.CalledProcessError as exc:
|
|
pytest.fail(exc.stderr.decode("utf8"))
|
|
|
|
test = pd.read_table(tester.output)
|
|
baseline = pd.read_table(tester.baseline_file)
|
|
|
|
test = test.reindex(columns=sorted(test.columns))
|
|
assert_frame_equal(test, baseline, check_like=True)
|
|
|
|
def test_diff():
|
|
tester = WikiqTester(SAILORMOON, "diff", in_compression="7z", out_format='jsonl')
|
|
|
|
try:
|
|
tester.call_wikiq("--diff", "--fandom-2020")
|
|
except subprocess.CalledProcessError as exc:
|
|
pytest.fail(exc.stderr.decode("utf8"))
|
|
|
|
test = pd.read_json(tester.output, lines=True)
|
|
assert "diff" in test.columns, "diff column should exist"
|
|
assert "diff_timeout" in test.columns, "diff_timeout column should exist"
|
|
assert len(test) > 0, "Should have output rows"
|
|
|
|
def test_diff_plus_pwr():
|
|
tester = WikiqTester(SAILORMOON, "diff_pwr", in_compression="7z", out_format='jsonl')
|
|
|
|
try:
|
|
tester.call_wikiq("--diff --persistence wikidiff2", "--fandom-2020")
|
|
except subprocess.CalledProcessError as exc:
|
|
pytest.fail(exc.stderr.decode("utf8"))
|
|
|
|
test = pd.read_json(tester.output, lines=True)
|
|
assert "diff" in test.columns, "diff column should exist"
|
|
assert "token_revs" in test.columns, "token_revs column should exist"
|
|
assert len(test) > 0, "Should have output rows"
|
|
|
|
def test_text():
|
|
tester = WikiqTester(SAILORMOON, "text", in_compression="7z", out_format='jsonl')
|
|
|
|
try:
|
|
tester.call_wikiq("--diff", "--text","--fandom-2020")
|
|
except subprocess.CalledProcessError as exc:
|
|
pytest.fail(exc.stderr.decode("utf8"))
|
|
|
|
test = pd.read_json(tester.output, lines=True)
|
|
assert "text" in test.columns, "text column should exist"
|
|
assert "diff" in test.columns, "diff column should exist"
|
|
assert len(test) > 0, "Should have output rows"
|
|
|
|
def test_malformed_noargs():
|
|
tester = WikiqTester(wiki=TWINPEAKS, case_name="noargs", in_compression="7z")
|
|
want_exception = (
|
|
"xml.etree.ElementTree.ParseError: no element found: line 1369, column 0"
|
|
)
|
|
|
|
try:
|
|
tester.call_wikiq()
|
|
except subprocess.CalledProcessError as exc:
|
|
errlines = exc.stderr.decode("utf8").splitlines()
|
|
assert errlines[-1] == want_exception
|
|
else:
|
|
pytest.fail("No exception raised, want: {}".format(want_exception))
|
|
|
|
def test_stdout_noargs():
|
|
tester = WikiqTester(wiki=SAILORMOON, case_name="noargs", in_compression="7z")
|
|
|
|
try:
|
|
outs = tester.call_wikiq("--stdout", "--fandom-2020", out=False).decode(
|
|
"utf8"
|
|
)
|
|
except subprocess.CalledProcessError as exc:
|
|
pytest.fail(exc.stderr.decode("utf8"))
|
|
|
|
test = pd.read_table(StringIO(outs))
|
|
baseline = pd.read_table(tester.baseline_file)
|
|
assert_frame_equal(test, baseline, check_like=True)
|
|
|
|
def test_bad_regex():
|
|
tester = WikiqTester(wiki=REGEXTEST, case_name="bad_regex")
|
|
|
|
# sample arguments for checking that bad arguments get terminated / test_regex_arguments
|
|
bad_arguments_list = [
|
|
# label is missing
|
|
"-RP '\\b\\d+\\b'",
|
|
# number of reg and number of labels do not match
|
|
"-RP 'NPO V' -RP THE -RPl testlabel",
|
|
# cp but rp label
|
|
"-CP '(Tamil|Li)' -RPl testlabel",
|
|
# regex is missing
|
|
"-CPl testlabel",
|
|
"-RP '\\b\\w{3}\\b' -RPl threeletters -CP '\\b\\w{3}\\b'",
|
|
]
|
|
|
|
for arguments in bad_arguments_list:
|
|
try:
|
|
tester.call_wikiq("--stdout", arguments, out=False)
|
|
except subprocess.CalledProcessError as exc:
|
|
# we want to check that the bad arguments were caught and sys.exit is stopping the code
|
|
print(exc.stderr.decode("utf-8"))
|
|
else:
|
|
pytest.fail("No exception raised, want Exception")
|
|
|
|
def test_good_regex():
|
|
# sample arguments for checking the outcomes of good arguments / test_basic_regex
|
|
good_arguments_list = [
|
|
"-RP '\\b\\d{3}\\b' -RPl threedigits",
|
|
"-RP 'TestCase' -RP 'page' -RPl testcases -RPl page_word",
|
|
"-CP 'Chevalier' -CPl chev_com -RP 'welcome to Wikipedia' -RPl wiki_welcome -CP 'Warning' -CPl warning",
|
|
"-CP 'WP:EVADE' -CPl wp_evade",
|
|
]
|
|
|
|
for i, arguments in enumerate(good_arguments_list):
|
|
tester = WikiqTester(wiki=REGEXTEST, case_name="basic", suffix=str(i))
|
|
|
|
try:
|
|
tester.call_wikiq(arguments)
|
|
except subprocess.CalledProcessError as exc:
|
|
pytest.fail(exc.stderr.decode("utf8"))
|
|
|
|
test = pd.read_table(tester.output)
|
|
|
|
baseline = pd.read_table(tester.baseline_file)
|
|
assert_frame_equal(test, baseline, check_like=True)
|
|
print(i)
|
|
|
|
def test_capturegroup_regex():
|
|
cap_arguments_list = [
|
|
"-RP 'Li Chevalier' -RPl li_cheval -CP '(?P<letter>\\b[a-zA-Z]{3}\\b)|(?P<number>\\b\\d+\\b)|(?P<cat>\\bcat\\b)' -CPl three",
|
|
"-CP '(?P<a>\\bTestCaseA\\b)|(?P<b>\\bTestCaseB\\b)|(?P<c>\\bTestCaseC\\b)|(?P<d>\\bTestCaseD\\b)' -CPl testcase -RP '(?P<npov>npov|NPOV)|(?P<neutral>neutral point of view)' -RPl npov",
|
|
]
|
|
|
|
for i, arguments in enumerate(cap_arguments_list):
|
|
tester = WikiqTester(
|
|
wiki=REGEXTEST, case_name="capturegroup", suffix=str(i)
|
|
)
|
|
|
|
try:
|
|
tester.call_wikiq(arguments)
|
|
except subprocess.CalledProcessError as exc:
|
|
pytest.fail(exc.stderr.decode("utf8"))
|
|
|
|
test = pd.read_table(tester.output)
|
|
|
|
baseline = pd.read_table(tester.baseline_file)
|
|
assert_frame_equal(test, baseline, check_like=True)
|
|
|
|
def test_regex_none_content():
|
|
# deleted or suppressed revisions yield None for text and comments, and
|
|
# matchmake must tolerate that in both the capture-group and plain paths
|
|
pair = RegexPair(r"(?P<letter>\b[a-zA-Z]{3}\b)|(?P<number>\b\d+\b)", "cap")
|
|
assert pair.matchmake(None) == {"cap_letter": None, "cap_number": None}
|
|
|
|
pair = RegexPair(r"\b\d{3}\b", "digits")
|
|
assert pair.matchmake(None) == {"digits": None}
|
|
|
|
def test_regex_deleted_revisions():
|
|
# the ikwiki dump contains revisions with deleted text and deleted
|
|
# comments; regex matching must handle them rather than crashing
|
|
tester = WikiqTester(wiki=IKWIKI, case_name="regex_deleted")
|
|
|
|
try:
|
|
tester.call_wikiq(
|
|
"-RP '(?P<npov>npov|NPOV)' -RPl npov",
|
|
"-CP '(?P<talk>[Tt]alk)' -CPl talk",
|
|
)
|
|
except subprocess.CalledProcessError as exc:
|
|
pytest.fail(exc.stderr.decode("utf8"))
|
|
|
|
test = pd.read_table(tester.output)
|
|
|
|
deleted = test[test["deleted"]]
|
|
assert len(deleted) > 0
|
|
assert deleted["npov_npov"].isna().all()
|
|
assert deleted["talk_talk"].isna().all()
|
|
|
|
def test_regex_count_only():
|
|
# count mode returns match counts instead of the matched text, with 0
|
|
# (never None) for revisions with no matches or deleted content
|
|
pair = RegexPair(r"\bcat\b", "cats", count_only=True)
|
|
assert pair.matchmake("cat cat dog cat") == {"cats": 3}
|
|
assert pair.matchmake("dog") == {"cats": 0}
|
|
assert pair.matchmake(None) == {"cats": 0}
|
|
assert pair.get_pyarrow_fields()[0].type == pa.int64()
|
|
|
|
pair = RegexPair(r"(?P<a>\bcat\b)|(?P<b>\bdog\b)", "pets", count_only=True)
|
|
assert pair.matchmake("cat dog cat") == {"pets_a": 2, "pets_b": 1}
|
|
assert pair.matchmake("bird") == {"pets_a": 0, "pets_b": 0}
|
|
assert pair.matchmake(None) == {"pets_a": 0, "pets_b": 0}
|
|
assert all(f.type == pa.int64() for f in pair.get_pyarrow_fields())
|
|
|
|
# without count_only the original behavior is unchanged
|
|
pair = RegexPair(r"\bcat\b", "cats")
|
|
assert pair.matchmake("cat cat dog cat") == {"cats": "cat, cat, cat"}
|
|
assert pair.get_pyarrow_fields()[0].type == pa.string()
|
|
|
|
def test_regex_count_e2e():
|
|
# -RPc/-CPc produce integer columns whose values equal the number of
|
|
# matches in the revision text and comment respectively
|
|
tester = WikiqTester(wiki=REGEXTEST, case_name="regex_count")
|
|
|
|
try:
|
|
tester.call_wikiq(
|
|
"--text",
|
|
"-RP 'TestCase' -RPl testcases -RPc",
|
|
"-CP '\\b[Cc]hevalier\\b' -CPl chev -CPc",
|
|
)
|
|
except subprocess.CalledProcessError as exc:
|
|
pytest.fail(exc.stderr.decode("utf8"))
|
|
|
|
import re
|
|
test = pd.read_table(tester.output)
|
|
|
|
assert test["testcases"].dtype == "int64"
|
|
assert test["chev"].dtype == "int64"
|
|
|
|
for _, row in test.iterrows():
|
|
text = row["text"] if isinstance(row["text"], str) else None
|
|
expected = len(re.findall("TestCase", text)) if text is not None else 0
|
|
assert row["testcases"] == expected, f"revid {row['revid']}: {row['testcases']} != {expected}"
|
|
|
|
comment = row["edit_summary"] if isinstance(row["edit_summary"], str) else None
|
|
expected = len(re.findall(r"\b[Cc]hevalier\b", comment)) if comment is not None else 0
|
|
assert row["chev"] == expected, f"revid {row['revid']}: {row['chev']} != {expected}"
|
|
|
|
def test_content_sizes():
|
|
# hand-computed fixtures: (text, prose_chars, structured_chars)
|
|
cases = [
|
|
# pure prose
|
|
("Hello world.", 12, 0),
|
|
# prose plus a listing-style template: parameter values count as
|
|
# structured, parameter names and braces do not
|
|
# prose: "Visit us. " = 10; structured: Hotel(5) + 555(3) + Nice place(10) = 18
|
|
("Visit us. {{listing|name=Hotel|phone=555|description=Nice place}}", 10, 18),
|
|
# tag inner text counts, attributes do not
|
|
# prose: "Before "(7) + " after"(6) = 13; structured: "inner text" = 10
|
|
('Before <span class="x">inner text</span> after', 13, 10),
|
|
# nesting: a template inside a template argument stays structured
|
|
# structured: deep(4) + " shallow"(8) = 12
|
|
("{{outer|a={{inner|b=deep}} shallow}}", 0, 12),
|
|
# a labeled wikilink in prose counts its label; a wikilink inside a
|
|
# template argument counts its rendered text as structured
|
|
# prose: "See "(4) + "the main page"(13) + " and "(5) = 22; structured: Help(4)
|
|
("See [[Main Page|the main page]] and {{box|link=[[Help]]}}", 22, 4),
|
|
# HTML comments count in neither column
|
|
# prose: "Text "(5) + " more"(5) = 10
|
|
("Text <!-- hidden --> more", 10, 0),
|
|
]
|
|
|
|
parser = WikitextParser()
|
|
for text, prose, structured in cases:
|
|
result = parser.content_sizes(text)
|
|
assert result == (prose, structured), f"{text!r}: got {result}, want {(prose, structured)}"
|
|
assert result[0] + result[1] <= len(text), f"{text!r}: content exceeds total size"
|
|
|
|
assert parser.content_sizes(None) is None
|
|
|
|
def test_content_sizes_e2e():
|
|
# the ikwiki dump contains revisions with deleted text, which must get
|
|
# null content sizes; all other revisions must satisfy the invariant
|
|
# prose_chars + structured_chars <= len(text)
|
|
tester = WikiqTester(IKWIKI, "content_sizes")
|
|
|
|
try:
|
|
tester.call_wikiq("--content-sizes", "--text")
|
|
except subprocess.CalledProcessError as exc:
|
|
pytest.fail(exc.stderr.decode("utf8"))
|
|
|
|
test = pd.read_table(tester.output)
|
|
|
|
assert "prose_chars" in test.columns
|
|
assert "structured_chars" in test.columns
|
|
|
|
deleted = test[test["deleted"]]
|
|
assert len(deleted) > 0
|
|
assert deleted["prose_chars"].isna().all()
|
|
assert deleted["structured_chars"].isna().all()
|
|
|
|
present = test[~test["deleted"] & test["text"].notna() & test["prose_chars"].notna()]
|
|
assert len(present) > 0
|
|
total = present["prose_chars"] + present["structured_chars"]
|
|
assert (total <= present["text"].str.len()).all()
|
|
|
|
def test_external_links_only():
|
|
"""Test that --external-links extracts external links correctly."""
|
|
import mwparserfromhell
|
|
|
|
tester = WikiqTester(SAILORMOON, "external_links_only", in_compression="7z", out_format="jsonl")
|
|
|
|
try:
|
|
# Also include --text so we can verify extraction against actual wikitext
|
|
tester.call_wikiq("--external-links", "--text", "--fandom-2020")
|
|
except subprocess.CalledProcessError as exc:
|
|
pytest.fail(exc.stderr.decode("utf8"))
|
|
|
|
test = pd.read_json(tester.output, lines=True)
|
|
|
|
# Verify external_links column exists
|
|
assert "external_links" in test.columns, "external_links column should exist"
|
|
|
|
# Verify citations column does NOT exist
|
|
assert "citations" not in test.columns, "citations column should NOT exist when only --external-links is used"
|
|
|
|
# Verify column has list/array type (pandas reads pyarrow lists as numpy arrays)
|
|
assert test["external_links"].apply(lambda x: x is None or hasattr(x, '__len__')).all(), \
|
|
"external_links should be a list/array type or None"
|
|
|
|
# Verify that extracted URLs look like valid URIs (have a scheme or are protocol-relative)
|
|
all_urls = []
|
|
for links in test["external_links"]:
|
|
if links is not None and len(links) > 0:
|
|
all_urls.extend(links)
|
|
|
|
for url in all_urls:
|
|
# External links can be http, https, mailto, ftp, etc. or protocol-relative (//)
|
|
has_scheme = ":" in url and url.index(":") < 10 # scheme:... with short scheme
|
|
is_protocol_relative = url.startswith("//")
|
|
assert has_scheme or is_protocol_relative, \
|
|
f"External link should be a valid URI, got: {url}"
|
|
|
|
# Verify extraction matches mwparserfromhell for a sample of rows with text
|
|
rows_with_links = test[test["external_links"].apply(lambda x: x is not None and len(x) > 0)]
|
|
if len(rows_with_links) > 0:
|
|
# Test up to 5 rows
|
|
sample = rows_with_links.head(5)
|
|
for idx, row in sample.iterrows():
|
|
text = row["text"]
|
|
if text:
|
|
wikicode = mwparserfromhell.parse(text)
|
|
expected_links = [str(link.url) for link in wikicode.filter_external_links()]
|
|
actual_links = list(row["external_links"])
|
|
assert actual_links == expected_links, \
|
|
f"Row {idx}: external_links mismatch. Expected {expected_links}, got {actual_links}"
|
|
|
|
print(f"External links only test passed! {len(test)} rows, {len(all_urls)} total URLs extracted")
|
|
|
|
|
|
def test_citations_only():
|
|
"""Test that --citations extracts citations correctly."""
|
|
import mwparserfromhell
|
|
from wikiq.wikitext_parser import WikitextParser
|
|
|
|
tester = WikiqTester(SAILORMOON, "citations_only", in_compression="7z", out_format="jsonl")
|
|
|
|
try:
|
|
# Also include --text so we can verify extraction against actual wikitext
|
|
tester.call_wikiq("--citations", "--text", "--fandom-2020")
|
|
except subprocess.CalledProcessError as exc:
|
|
pytest.fail(exc.stderr.decode("utf8"))
|
|
|
|
test = pd.read_json(tester.output, lines=True)
|
|
|
|
# Verify citations column exists
|
|
assert "citations" in test.columns, "citations column should exist"
|
|
|
|
# Verify external_links column does NOT exist
|
|
assert "external_links" not in test.columns, "external_links column should NOT exist when only --citations is used"
|
|
|
|
# Verify column has list/array type (pandas reads pyarrow lists as numpy arrays)
|
|
assert test["citations"].apply(lambda x: x is None or hasattr(x, '__len__')).all(), \
|
|
"citations should be a list/array type or None"
|
|
|
|
# Verify that extracted citations have correct prefixes (ref: or template:)
|
|
all_citations = []
|
|
for citations in test["citations"]:
|
|
if citations is not None and len(citations) > 0:
|
|
all_citations.extend(citations)
|
|
|
|
for citation in all_citations:
|
|
assert citation.startswith("ref:") or citation.startswith("template:"), \
|
|
f"Citation should start with 'ref:' or 'template:', got: {citation}"
|
|
|
|
# Verify extraction matches WikitextParser for a sample of rows with text
|
|
rows_with_citations = test[test["citations"].apply(lambda x: x is not None and len(x) > 0)]
|
|
if len(rows_with_citations) > 0:
|
|
parser = WikitextParser()
|
|
# Test up to 5 rows
|
|
sample = rows_with_citations.head(5)
|
|
for idx, row in sample.iterrows():
|
|
text = row["text"]
|
|
if text:
|
|
expected_citations = parser.extract_citations(text)
|
|
actual_citations = list(row["citations"])
|
|
assert actual_citations == expected_citations, \
|
|
f"Row {idx}: citations mismatch. Expected {expected_citations}, got {actual_citations}"
|
|
|
|
print(f"Citations only test passed! {len(test)} rows, {len(all_citations)} total citations extracted")
|
|
|
|
|
|
def test_external_links_and_citations():
|
|
"""Test that both --external-links and --citations work together (shared parser)."""
|
|
import mwparserfromhell
|
|
from wikiq.wikitext_parser import WikitextParser
|
|
|
|
tester = WikiqTester(SAILORMOON, "external_links_and_citations", in_compression="7z", out_format="jsonl")
|
|
|
|
try:
|
|
# Also include --text so we can verify extraction against actual wikitext
|
|
tester.call_wikiq("--external-links", "--citations", "--text", "--fandom-2020")
|
|
except subprocess.CalledProcessError as exc:
|
|
pytest.fail(exc.stderr.decode("utf8"))
|
|
|
|
test = pd.read_json(tester.output, lines=True)
|
|
|
|
# Verify both columns exist
|
|
assert "external_links" in test.columns, "external_links column should exist"
|
|
assert "citations" in test.columns, "citations column should exist"
|
|
|
|
# Verify both columns have list/array types (pandas reads pyarrow lists as numpy arrays)
|
|
assert test["external_links"].apply(lambda x: x is None or hasattr(x, '__len__')).all(), \
|
|
"external_links should be a list/array type or None"
|
|
assert test["citations"].apply(lambda x: x is None or hasattr(x, '__len__')).all(), \
|
|
"citations should be a list/array type or None"
|
|
|
|
# Verify URLs look like valid URIs (have a scheme or are protocol-relative)
|
|
all_urls = []
|
|
for links in test["external_links"]:
|
|
if links is not None and len(links) > 0:
|
|
all_urls.extend(links)
|
|
|
|
for url in all_urls:
|
|
# External links can be http, https, mailto, ftp, etc. or protocol-relative (//)
|
|
has_scheme = ":" in url and url.index(":") < 10 # scheme:... with short scheme
|
|
is_protocol_relative = url.startswith("//")
|
|
assert has_scheme or is_protocol_relative, \
|
|
f"External link should be a valid URI, got: {url}"
|
|
|
|
# Verify citations have correct prefixes
|
|
all_citations = []
|
|
for citations in test["citations"]:
|
|
if citations is not None and len(citations) > 0:
|
|
all_citations.extend(citations)
|
|
|
|
for citation in all_citations:
|
|
assert citation.startswith("ref:") or citation.startswith("template:"), \
|
|
f"Citation should start with 'ref:' or 'template:', got: {citation}"
|
|
|
|
# Verify extraction matches WikitextParser for a sample of rows with text
|
|
# This tests that the shared parser optimization works correctly
|
|
parser = WikitextParser()
|
|
rows_with_content = test[
|
|
(test["external_links"].apply(lambda x: x is not None and len(x) > 0)) |
|
|
(test["citations"].apply(lambda x: x is not None and len(x) > 0))
|
|
]
|
|
if len(rows_with_content) > 0:
|
|
# Test up to 5 rows
|
|
sample = rows_with_content.head(5)
|
|
for idx, row in sample.iterrows():
|
|
text = row["text"]
|
|
if text:
|
|
# Verify external links
|
|
wikicode = mwparserfromhell.parse(text)
|
|
expected_links = [str(link.url) for link in wikicode.filter_external_links()]
|
|
actual_links = list(row["external_links"]) if row["external_links"] is not None else []
|
|
assert actual_links == expected_links, \
|
|
f"Row {idx}: external_links mismatch. Expected {expected_links}, got {actual_links}"
|
|
|
|
# Verify citations
|
|
expected_citations = parser.extract_citations(text)
|
|
actual_citations = list(row["citations"]) if row["citations"] is not None else []
|
|
assert actual_citations == expected_citations, \
|
|
f"Row {idx}: citations mismatch. Expected {expected_citations}, got {actual_citations}"
|
|
|
|
print(f"External links and citations test passed! {len(test)} rows, {len(all_urls)} URLs, {len(all_citations)} citations")
|
|
|
|
|
|
def test_no_wikitext_columns():
|
|
"""Test that neither external_links nor citations columns exist without flags."""
|
|
tester = WikiqTester(SAILORMOON, "no_wikitext_columns", in_compression="7z", out_format="jsonl")
|
|
|
|
try:
|
|
tester.call_wikiq("--fandom-2020")
|
|
except subprocess.CalledProcessError as exc:
|
|
pytest.fail(exc.stderr.decode("utf8"))
|
|
|
|
test = pd.read_json(tester.output, lines=True)
|
|
|
|
# Verify neither column exists
|
|
assert "external_links" not in test.columns, "external_links column should NOT exist without --external-links flag"
|
|
assert "citations" not in test.columns, "citations column should NOT exist without --citations flag"
|
|
|
|
print(f"No wikitext columns test passed! {len(test)} rows processed")
|
|
|
|
|
|
def test_wikilinks():
|
|
"""Test that --wikilinks extracts internal wikilinks correctly."""
|
|
import mwparserfromhell
|
|
|
|
tester = WikiqTester(SAILORMOON, "wikilinks", in_compression="7z", out_format="jsonl")
|
|
|
|
try:
|
|
tester.call_wikiq("--wikilinks", "--text", "--fandom-2020")
|
|
except subprocess.CalledProcessError as exc:
|
|
pytest.fail(exc.stderr.decode("utf8"))
|
|
|
|
test = pd.read_json(tester.output, lines=True)
|
|
|
|
# Verify wikilinks column exists
|
|
assert "wikilinks" in test.columns, "wikilinks column should exist"
|
|
|
|
# Verify column has list/array type
|
|
assert test["wikilinks"].apply(lambda x: x is None or hasattr(x, '__len__')).all()
|
|
|
|
# Verify extraction matches mwparserfromhell for sample rows
|
|
rows_with_links = test[test["wikilinks"].apply(lambda x: x is not None and len(x) > 0)]
|
|
if len(rows_with_links) > 0:
|
|
sample = rows_with_links.head(5)
|
|
for idx, row in sample.iterrows():
|
|
text = row["text"]
|
|
if text:
|
|
wikicode = mwparserfromhell.parse(text)
|
|
expected = []
|
|
for link in wikicode.filter_wikilinks():
|
|
title = str(link.title).strip()
|
|
display_text = str(link.text).strip() if link.text else None
|
|
expected.append({"title": title, "text": display_text})
|
|
|
|
actual = list(row["wikilinks"])
|
|
# Convert to comparable format (pandas may read as dicts or named tuples)
|
|
actual_dicts = [{"title": item["title"], "text": item["text"]} for item in actual]
|
|
assert actual_dicts == expected, f"Row {idx}: wikilinks mismatch"
|
|
|
|
print(f"Wikilinks test passed! {len(test)} rows processed")
|
|
|
|
|
|
def test_templates():
|
|
"""Test that --templates extracts templates correctly."""
|
|
import mwparserfromhell
|
|
|
|
tester = WikiqTester(SAILORMOON, "templates", in_compression="7z", out_format="jsonl")
|
|
|
|
try:
|
|
tester.call_wikiq("--templates", "--text", "--fandom-2020")
|
|
except subprocess.CalledProcessError as exc:
|
|
pytest.fail(exc.stderr.decode("utf8"))
|
|
|
|
test = pd.read_json(tester.output, lines=True)
|
|
|
|
# Verify templates column exists
|
|
assert "templates" in test.columns, "templates column should exist"
|
|
|
|
# Verify column has list/array type
|
|
assert test["templates"].apply(lambda x: x is None or hasattr(x, '__len__')).all()
|
|
|
|
# Verify extraction matches mwparserfromhell for sample rows
|
|
rows_with_templates = test[test["templates"].apply(lambda x: x is not None and len(x) > 0)]
|
|
if len(rows_with_templates) > 0:
|
|
sample = rows_with_templates.head(5)
|
|
for idx, row in sample.iterrows():
|
|
text = row["text"]
|
|
if text:
|
|
wikicode = mwparserfromhell.parse(text)
|
|
expected = []
|
|
for template in wikicode.filter_templates():
|
|
name = str(template.name).strip()
|
|
params = {}
|
|
for param in template.params:
|
|
param_name = str(param.name).strip()
|
|
param_value = str(param.value).strip()
|
|
params[param_name] = param_value
|
|
expected.append({"name": name, "params": params})
|
|
|
|
actual = list(row["templates"])
|
|
# Convert to comparable format
|
|
actual_list = []
|
|
for item in actual:
|
|
actual_list.append({
|
|
"name": item["name"],
|
|
"params": dict(item["params"]) if item["params"] else {}
|
|
})
|
|
assert actual_list == expected, f"Row {idx}: templates mismatch"
|
|
|
|
print(f"Templates test passed! {len(test)} rows processed")
|
|
|
|
|
|
def test_headings():
|
|
"""Test that --headings extracts section headings correctly."""
|
|
import mwparserfromhell
|
|
|
|
tester = WikiqTester(SAILORMOON, "headings", in_compression="7z", out_format="jsonl")
|
|
|
|
try:
|
|
tester.call_wikiq("--headings", "--text", "--fandom-2020")
|
|
except subprocess.CalledProcessError as exc:
|
|
pytest.fail(exc.stderr.decode("utf8"))
|
|
|
|
test = pd.read_json(tester.output, lines=True)
|
|
|
|
# Verify headings column exists
|
|
assert "headings" in test.columns, "headings column should exist"
|
|
|
|
# Verify column has list/array type
|
|
assert test["headings"].apply(lambda x: x is None or hasattr(x, '__len__')).all()
|
|
|
|
# Verify extraction matches mwparserfromhell for sample rows
|
|
rows_with_headings = test[test["headings"].apply(lambda x: x is not None and len(x) > 0)]
|
|
if len(rows_with_headings) > 0:
|
|
sample = rows_with_headings.head(5)
|
|
for idx, row in sample.iterrows():
|
|
text = row["text"]
|
|
if text:
|
|
wikicode = mwparserfromhell.parse(text)
|
|
expected = []
|
|
for heading in wikicode.filter_headings():
|
|
level = heading.level
|
|
heading_text = str(heading.title).strip()
|
|
expected.append({"level": level, "text": heading_text})
|
|
|
|
actual = list(row["headings"])
|
|
# Convert to comparable format
|
|
actual_list = [{"level": item["level"], "text": item["text"]} for item in actual]
|
|
assert actual_list == expected, f"Row {idx}: headings mismatch"
|
|
|
|
print(f"Headings test passed! {len(test)} rows processed")
|
|
|
|
|
|
|