add (optional) diff and text columns to output.

This commit is contained in:
Nathan TeBlunthuis
2025-07-07 14:39:52 -07:00
parent a8e9e7f4fd
commit d6c4c0a416
4 changed files with 373 additions and 295 deletions

48
wikiq
View File

@@ -208,6 +208,8 @@ class WikiqParser:
regex_match_comment: list[str],
regex_revision_label: list[str],
regex_comment_label: list[str],
text: bool = False,
diff: bool = False,
collapse_user: bool = False,
persist: int = None,
namespaces: Union[list[int], None] = None,
@@ -226,7 +228,8 @@ class WikiqParser:
self.persist: int = persist
self.namespaces = []
self.revert_radius = revert_radius
self.diff = diff
self.text = text
if namespaces is not None:
self.namespace_filter = set(namespaces)
else:
@@ -331,6 +334,9 @@ class WikiqParser:
tables.RevisionIsAnon(),
])
if self.text:
table.columns.append(tables.RevisionText())
if self.collapse_user:
table.columns.append(tables.RevisionCollapsed())
@@ -345,6 +351,10 @@ class WikiqParser:
schema = table.schema()
schema = schema.append(pa.field('revert', pa.bool_(), nullable=True))
if self.diff:
from diff_pyarrow_schema import diff_field
schema = schema.append(diff_field)
# Add regex fields to the schema.
for pair in self.regex_revision_pairs:
for field in pair.get_pyarrow_fields():
@@ -412,9 +422,18 @@ class WikiqParser:
revision_texts.append(rev.text)
wikidiff_matcher = None
if self.diff or self.persist == PersistMethod.wikidiff2:
wikidiff_matcher = WikiDiffMatcher(revision_texts,
tokenizer=wikitext_split,
)
# Collect the set of pages currently buffered in the table so we can run multi-page functions on them.
row_buffer = table.pop()
if self.diff:
row_buffer['diff'] = [[entry for entry in wikidiff_matcher.diffs[i]['diff'] if entry['type'] != 0 ] for i in range(len(revision_texts))]
is_revert_column: list[Union[bool, None]] = []
for r, d in zip(row_buffer['reverteds'], row_buffer['deleted']):
if self.revert_radius == 0 or d:
@@ -428,6 +447,7 @@ class WikiqParser:
row_buffer[k] = v
regex_matches = {}
if self.persist != PersistMethod.none:
window = deque(maxlen=PERSISTENCE_RADIUS)
@@ -443,9 +463,8 @@ class WikiqParser:
state = mwpersistence.DiffState(SegmentMatcher(tokenizer=wikitext_split),
revert_radius=PERSISTENCE_RADIUS)
elif self.persist == PersistMethod.wikidiff2:
state = mwpersistence.DiffState(WikiDiffMatcher(revision_texts,
tokenizer=wikitext_split,
),
state = mwpersistence.DiffState(wikidiff_matcher,
revert_radius=PERSISTENCE_RADIUS)
else:
from mw.lib import persistence
@@ -469,8 +488,6 @@ class WikiqParser:
row_buffer['tokens_removed'].append(len(old_tokens_removed))
row_buffer['tokens_window'].append(PERSISTENCE_RADIUS - 1)
del row_buffer['text']
# print out metadata for the last RADIUS revisions
for i, item in enumerate(window):
# if the window was full, we've already printed item 0
@@ -485,6 +502,9 @@ class WikiqParser:
row_buffer['tokens_removed'].append(len(tokens_removed))
row_buffer['tokens_window'].append(len(window) - (i + 1))
if not self.text:
del row_buffer['text']
writer.write(pa.table(row_buffer, schema=schema))
page_count += 1
@@ -494,7 +514,6 @@ class WikiqParser:
writer.close()
def match_archive_suffix(input_filename):
if re.match(r'.*\.7z$', input_filename):
cmd = ["7za", "x", "-so", input_filename]
@@ -580,6 +599,14 @@ def main():
action='append',
help="The label for the outputted column based on matching the regex in comments.")
parser.add_argument('-d', '--diff', dest="diff", default=False,
action='store_true',
help="Output a diff structure for each revision with information about changed or moved lines.")
parser.add_argument('-t', '--text', dest="text", default=False,
action='store_true',
help="Output the text of the revision.")
parser.add_argument('--fandom-2020', dest="fandom_2020",
action='store_true',
help="Whether the archive is from the fandom 2020 dumps by Wikiteam. These dumps can have multiple .xml files in their archives.")
@@ -604,6 +631,7 @@ def main():
else:
namespaces = None
print(args, file=sys.stderr)
if len(args.dumpfiles) > 0:
for filename in args.dumpfiles:
input_file = open_input_file(filename, args.fandom_2020)
@@ -637,6 +665,8 @@ def main():
regex_revision_label=args.regex_revision_label,
regex_match_comment=args.regex_match_comment,
regex_comment_label=args.regex_comment_label,
text=args.text,
diff=args.diff,
output_parquet=output_parquet,
)
@@ -656,7 +686,9 @@ def main():
regex_match_revision=args.regex_match_revision,
regex_revision_label=args.regex_revision_label,
regex_match_comment=args.regex_match_comment,
regex_comment_label=args.regex_comment_label)
regex_comment_label=args.regex_comment_label,
diff=args.diff,
text=args.text)
wikiq.process()