diff --git a/src/wikiq/__init__.py b/src/wikiq/__init__.py index 4ea2e50..9d3cb7a 100755 --- a/src/wikiq/__init__.py +++ b/src/wikiq/__init__.py @@ -207,10 +207,10 @@ def build_schema( return schema -def make_regex_pairs(patterns, labels) -> list: +def make_regex_pairs(patterns, labels, count_only=False) -> list: """Create RegexPair objects from patterns and labels.""" if (patterns is not None and labels is not None) and (len(patterns) == len(labels)): - return [RegexPair(pattern, label) for pattern, label in zip(patterns, labels)] + return [RegexPair(pattern, label, count_only) for pattern, label in zip(patterns, labels)] elif patterns is None and labels is None: return [] else: @@ -431,21 +431,23 @@ If the pattern does not include a capture group, then only one output column wil class RegexPair(object): - def __init__(self, pattern, label): + def __init__(self, pattern, label, count_only=False): self.pattern = re.compile(pattern) self.label = label + self.count_only = count_only self.has_groups = bool(self.pattern.groupindex) if self.has_groups: self.capture_groups = list(self.pattern.groupindex.keys()) def get_pyarrow_fields(self): + value_type = pa.int64() if self.count_only else pa.string() if self.has_groups: fields = [ - pa.field(self._make_key(cap_group), pa.string()) + pa.field(self._make_key(cap_group), value_type) for cap_group in self.capture_groups ] else: - fields = [pa.field(self.label, pa.string())] + fields = [pa.field(self.label, value_type)] return fields @@ -458,39 +460,50 @@ class RegexPair(object): if self.has_groups: # if there are matches of some sort in this revision content, fill the lists for each cap_group # content can be None when the text or comment was deleted/suppressed - if content is not None and self.pattern.search(content) is not None: - m = self.pattern.finditer(content) - matchobjects = list(m) - + if content is not None and (matchobjects := list(self.pattern.finditer(content))): for cap_group in self.capture_groups: key = self._make_key(cap_group) temp_list = [] for match in matchobjects: # we only want to add the match for the capture group if the match is not None - if match.group(cap_group) is not None: - temp_list.append(match.group(cap_group)) + if (group := match.group(cap_group)) is not None: + temp_list.append(group) - # if temp_list of matches is empty just make that column None + # if temp_list of matches is empty just make that column None (0 in count mode) if len(temp_list) == 0: - temp_dict[key] = None + if self.count_only: + temp_dict[key] = 0 + else: + temp_dict[key] = None # else we put in the list we made in the for-loop above else: - temp_dict[key] = ", ".join(temp_list) + if self.count_only: + temp_dict[key] = len(temp_list) + else: + temp_dict[key] = ", ".join(temp_list) - # there are no matches at all in this revision content, we default values to None + # there are no matches at all in this revision content, we default values to None (0 in count mode) else: for cap_group in self.capture_groups: key = self._make_key(cap_group) - temp_dict[key] = None + if self.count_only: + temp_dict[key] = 0 + else: + temp_dict[key] = None # there are no capture groups, we just search for all the matches of the regex else: # given that there are matches to be made - if content is not None and self.pattern.search(content) is not None: - m = self.pattern.findall(content) - temp_dict[self.label] = ", ".join(m) + if content is not None and (matches := self.pattern.findall(content)): + if self.count_only: + temp_dict[self.label] = len(matches) + else: + temp_dict[self.label] = ", ".join(matches) else: - temp_dict[self.label] = None + if self.count_only: + temp_dict[self.label] = 0 + else: + temp_dict[self.label] = None return temp_dict @@ -504,6 +517,8 @@ class WikiqParser: regex_match_comment: list[str], regex_revision_label: list[str], regex_comment_label: list[str], + regex_revision_output_count: bool = False, + regex_comment_output_count: bool = False, text: bool = False, diff: bool = False, collapse_user: bool = False, @@ -553,10 +568,10 @@ class WikiqParser: self.regex_schemas = [] self.regex_revision_pairs: list[RegexPair] = self.make_matchmake_pairs( - regex_match_revision, regex_revision_label + regex_match_revision, regex_revision_label, regex_revision_output_count ) self.regex_comment_pairs: list[RegexPair] = self.make_matchmake_pairs( - regex_match_comment, regex_comment_label + regex_match_comment, regex_comment_label, regex_comment_output_count ) # Initialize output @@ -599,13 +614,13 @@ class WikiqParser: if timer is not None: timer.cancel() - def make_matchmake_pairs(self, patterns, labels) -> list[RegexPair]: + def make_matchmake_pairs(self, patterns, labels, count_only=False) -> list[RegexPair]: if (patterns is not None and labels is not None) and ( len(patterns) == len(labels) ): result: list[RegexPair] = [] for pattern, label in zip(patterns, labels): - rp = RegexPair(pattern, label) + rp = RegexPair(pattern, label, count_only) result.append(rp) self.regex_schemas = self.regex_schemas + rp.get_pyarrow_fields() return result @@ -1137,6 +1152,14 @@ def main(): help="The label for the outputted column based on matching the regex in revision text.", ) + parser.add_argument( + "-RPc", + "--revision-pattern-count", + dest="regex_revision_output_count", + action="store_true", + help="If present, this will cause the revision patterns to return counts of the number of matches instead of the text of the matches themselves. It will affect all revision patterns.", + ) + parser.add_argument( "-CP", "--comment-pattern", @@ -1157,6 +1180,14 @@ def main(): help="The label for the outputted column based on matching the regex in comments.", ) + parser.add_argument( + "-CPc", + "--comment-pattern-count", + dest="regex_comment_output_count", + action="store_true", + help="If present, this will cause the comment patterns to return counts of the number of matches instead of the text of the matches themselves. It will affect all comment patterns.", + ) + parser.add_argument( "-d", "--diff", @@ -1267,8 +1298,8 @@ def main(): # Handle --print-schema: build and output schema, then exit if args.print_schema: - regex_revision_pairs = make_regex_pairs(args.regex_match_revision, args.regex_revision_label) - regex_comment_pairs = make_regex_pairs(args.regex_match_comment, args.regex_comment_label) + regex_revision_pairs = make_regex_pairs(args.regex_match_revision, args.regex_revision_label, args.regex_revision_output_count) + regex_comment_pairs = make_regex_pairs(args.regex_match_comment, args.regex_comment_label, args.regex_comment_output_count) table, _, _ = build_table( text=args.text, @@ -1344,8 +1375,10 @@ def main(): revert_radius=args.revert_radius, regex_match_revision=args.regex_match_revision, regex_revision_label=args.regex_revision_label, + regex_revision_output_count=args.regex_revision_output_count, regex_match_comment=args.regex_match_comment, regex_comment_label=args.regex_comment_label, + regex_comment_output_count=args.regex_comment_output_count, text=args.text, diff=args.diff, output_jsonl=output_jsonl, @@ -1400,8 +1433,10 @@ def main(): revert_radius=args.revert_radius, regex_match_revision=args.regex_match_revision, regex_revision_label=args.regex_revision_label, + regex_revision_output_count=args.regex_revision_output_count, regex_match_comment=args.regex_match_comment, regex_comment_label=args.regex_comment_label, + regex_comment_output_count=args.regex_comment_output_count, diff=args.diff, text=args.text, batch_size=args.batch_size, diff --git a/test/Wikiq_Unit_Test.py b/test/Wikiq_Unit_Test.py index e4d91aa..e987028 100644 --- a/test/Wikiq_Unit_Test.py +++ b/test/Wikiq_Unit_Test.py @@ -420,6 +420,55 @@ def test_regex_deleted_revisions(): 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\bcat\b)|(?P\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_external_links_only(): """Test that --external-links extracts external links correctly.""" import mwparserfromhell