add regex match counting (-RPc / -CPc)

Port the regex counting feature from the mako_changes-20230429 branch
(2ff4d60): the -RPc/--revision-pattern-count and
-CPc/--comment-pattern-count flags cause revision and comment patterns
to output the number of matches instead of the matched text, with 0
(never null) for revisions with no matches or deleted content. The
flags apply to all revision or comment patterns respectively, and count
columns are typed int64 in the output schema.

Counting the matches of common or large patterns previously required
returning the full matched text, which made for very large outputs.

Also take the code review suggestion from that branch (933ca75): use
assignment expressions so matching no longer makes redundant calls to
search() before finditer()/findall() or calls match.group() twice.

Includes unit tests for the count semantics and schema types, and an
end-to-end test that verifies counts against the number of matches
recomputed from each revision's text and comment.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
(cherry picked from commit 0b6e0fbc70)

src/wikiq/__init__.py

@@ -274,10 +274,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:
@@ -502,21 +502,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
 
@@ -529,37 +531,48 @@ 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:
+            if content is not None and (matchobjects := list(self.pattern.finditer(content))):
-                m = self.pattern.finditer(content)
-                matchobjects = list(m)
-
                 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:
+                        if (group := match.group(cap_group)) is not None:
-                            temp_list.append(match.group(cap_group))
+                            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:
+                        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:
+                        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)
+                    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:
+            if content is not None and (matches := self.pattern.findall(content)):
-                m = self.pattern.findall(content)
+                if self.count_only:
-                temp_dict[self.label] = ", ".join(m)
+                    temp_dict[self.label] = len(matches)
+                else:
+                    temp_dict[self.label] = ", ".join(matches)
+            else:
+                if self.count_only:
+                    temp_dict[self.label] = 0
                 else:
                     temp_dict[self.label] = None
 
@@ -575,6 +588,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,
@@ -626,10 +641,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
@@ -672,13 +687,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
@@ -1220,6 +1235,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",
@@ -1240,6 +1263,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",
@@ -1378,8 +1409,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_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)
+        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,
@@ -1457,8 +1488,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,
@@ -1514,8 +1547,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,

test/test_wikiq.py

@@ -427,6 +427,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<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_redirect_detection():
     from wikiq.tables import RedirectDetector