write regex captures to parquet arrays.

[mediawiki_dump_tools.git] / wikiq

diff --git a/wikiq b/wikiq
index a377d19a9469028a7763b37c9b2f46313c72b27c..be90a8b6b3bf8ef201f324731c73134162f757de 100755 (executable)
--- a/wikiq
+++ b/wikiq
@@ -9,6 +9,7 @@ import sys
 import os, os.path
 import re
 from datetime import datetime,timezone
 import os, os.path
 import re
 from datetime import datetime,timezone

+import json

 
 from subprocess import Popen, PIPE
 from collections import deque
 
 from subprocess import Popen, PIPE
 from collections import deque


@@ -29,6 +30,7 @@ import dataclasses as dc
 from dataclasses import dataclass
 import pyarrow as pa
 import pyarrow.parquet as pq
 from dataclasses import dataclass
 import pyarrow as pa
 import pyarrow.parquet as pq

+from itertools import chain

 
 class PersistMethod:
     none = 0
 
 class PersistMethod:
     none = 0


@@ -140,7 +142,11 @@ If the pattern does not include a capture group, then only one output column wil
 """
 class RegexPair(object):
     def __init__(self, pattern, label):
 """
 class RegexPair(object):
     def __init__(self, pattern, label):

-        self.pattern = re.compile(pattern)
+        self.pattern = pattern
+
+        if type(self.pattern) is str:
+            self.pattern = re.compile(pattern)
+

         self.label = label
         self.has_groups = bool(self.pattern.groupindex)
         if self.has_groups:
         self.label = label
         self.has_groups = bool(self.pattern.groupindex)
         if self.has_groups:


@@ -196,7 +202,7 @@ class RegexPair(object):
             if type(content) in(str, bytes):
                 if self.pattern.search(content) is not None:
                     m = self.pattern.findall(content)
             if type(content) in(str, bytes):
                 if self.pattern.search(content) is not None:
                     m = self.pattern.findall(content)

-                    temp_dict[self.label] = ', '.join(m)
+                    temp_dict[self.label] = m

                 else:
                     temp_dict[self.label] = None
 
                 else:
                     temp_dict[self.label] = None
 


@@ -254,13 +260,13 @@ class RevDataBase():
         pa.field("title",pa.string()),
         pa.field("namespace",pa.int32()),
         pa.field("deleted",pa.bool_()),
         pa.field("title",pa.string()),
         pa.field("namespace",pa.int32()),
         pa.field("deleted",pa.bool_()),

-        pa.field("test_chars",pa.int32()),
+        pa.field("text_chars",pa.int32()),

         pa.field("revert",pa.bool_()),
         pa.field("reverteds",pa.list_(pa.int64())),
         pa.field("sha1",pa.string()),
         pa.field("minor",pa.bool_()),
         pa.field("editor",pa.string()),
         pa.field("revert",pa.bool_()),
         pa.field("reverteds",pa.list_(pa.int64())),
         pa.field("sha1",pa.string()),
         pa.field("minor",pa.bool_()),
         pa.field("editor",pa.string()),

-        pa.field("anon",pa.bool_())
+        pa.field("anon",pa.bool_()),

     ]
 
     # pyarrow is a columnar format, so most of the work happens in the flush_parquet_buffer function
     ]
 
     # pyarrow is a columnar format, so most of the work happens in the flush_parquet_buffer function


@@ -291,6 +297,9 @@ class RevDataBase():
             elif f.type == list[int]:
                 row.append('"' + ",".join([str(x) for x in val]) + '"')
 
             elif f.type == list[int]:
                 row.append('"' + ",".join([str(x) for x in val]) + '"')
 

+            elif f.type == list[str]:
+                row.append('"' + ",".join([(x) for x in val]) + '"')
+

             elif f.type == str:
                 if self.urlencode and f.name in TO_ENCODE:
                     row.append(quote(str(val)))
             elif f.type == str:
                 if self.urlencode and f.name in TO_ENCODE:
                     row.append(quote(str(val)))


@@ -347,8 +356,10 @@ class RevDataCollapsePersistence uses multiple inheritence to make a class that
 class RevDataCollapsePersistence(RevDataCollapse, RevDataPersistence):
     pa_schema_fields = RevDataCollapse.pa_schema_fields + RevDataPersistence.pa_persistence_schema_fields
 
 class RevDataCollapsePersistence(RevDataCollapse, RevDataPersistence):
     pa_schema_fields = RevDataCollapse.pa_schema_fields + RevDataPersistence.pa_persistence_schema_fields
 

+
+

 class WikiqParser():
 class WikiqParser():

-    def __init__(self, input_file, output_file, regex_match_revision, regex_match_comment, regex_revision_label, regex_comment_label, collapse_user=False, persist=None, urlencode=False, namespaces = None, revert_radius=15, output_parquet=True, parquet_buffer_size=2000):
+    def __init__(self, input_file, output_file, regex_match_revision, regex_match_comment, regex_revision_label, regex_comment_label, collapse_user=False, persist=None, urlencode=False, namespaces = None, revert_radius=15, output_parquet=True, parquet_buffer_size=2000, siteinfo_file=None):

         """ 
         Parameters:
            persist : what persistence method to use. Takes a PersistMethod value
         """ 
         Parameters:
            persist : what persistence method to use. Takes a PersistMethod value


@@ -360,7 +371,7 @@ class WikiqParser():
         self.namespaces = []
         self.urlencode = urlencode
         self.revert_radius = revert_radius
         self.namespaces = []
         self.urlencode = urlencode
         self.revert_radius = revert_radius

-        
+

         if namespaces is not None:
             self.namespace_filter = set(namespaces)
         else:
         if namespaces is not None:
             self.namespace_filter = set(namespaces)
         else:


@@ -370,6 +381,24 @@ class WikiqParser():
         self.regex_revision_pairs = self.make_matchmake_pairs(regex_match_revision, regex_revision_label)
         self.regex_comment_pairs = self.make_matchmake_pairs(regex_match_comment, regex_comment_label)
 
         self.regex_revision_pairs = self.make_matchmake_pairs(regex_match_revision, regex_revision_label)
         self.regex_comment_pairs = self.make_matchmake_pairs(regex_match_comment, regex_comment_label)
 

+        if siteinfo_file is not None:
+            siteinfo = open_siteinfo(siteinfo_file)
+            siteinfo = json.loads(siteinfo.read())
+
+            magicwords = siteinfo.get('query').get('magicwords')
+        
+            if magicwords:
+                redirect_config = list(filter(lambda obj: obj.get("name") == "redirect", magicwords))
+                redirect_aliases = chain(* map(lambda obj: obj.get("aliases"), redirect_config))
+                redirect_aliases = list(map(lambda s: s.lstrip('#'), redirect_aliases))
+                redirect_aliases.append('REDIRECT') # just in case
+
+                # this regular expression is copied from pywikibot
+                pattern = '(?:' + '|'.join(redirect_aliases) + ')'
+                redirect_regex =  re.compile(r'\s*#{pattern}\s*:?\s*\[\[(.+?)(?:\|.*?)?\]\]'
+                                             .format(pattern=pattern), re.IGNORECASE | re.DOTALL)
+
+                self.regex_revision_pairs.extend(self.make_matchmake_pairs([redirect_regex], ["redirect"]))

 
         # This is where we set the type for revdata.
         
 
         # This is where we set the type for revdata.
         


@@ -477,7 +506,6 @@ class WikiqParser():
         page_count = 0
         rev_count = 0
 
         page_count = 0
         rev_count = 0
 

-

         # Iterate through pages
         for page in dump:
             namespace = page.namespace if page.namespace is not None else self.__get_namespace_from_title(page.title)
         # Iterate through pages
         for page in dump:
             namespace = page.namespace if page.namespace is not None else self.__get_namespace_from_title(page.title)


@@ -518,7 +546,7 @@ class WikiqParser():
                                              namespace = namespace
                                              )
 
                                              namespace = namespace
                                              )
 

-                rev_data = self.matchmake(rev, rev_data)
+                rev_data = self.matchmake_revision(rev, rev_data)

 
                 if not rev.deleted.text:
                     # rev.text can be None if the page has no text
 
                 if not rev.deleted.text:
                     # rev.text can be None if the page has no text


@@ -534,7 +562,7 @@ class WikiqParser():
                     rev_data.sha1 = text_sha1
 
                     # TODO rev.bytes doesn't work.. looks like a bug
                     rev_data.sha1 = text_sha1
 
                     # TODO rev.bytes doesn't work.. looks like a bug

-                    rev_data.text_chars = len(rev.text)
+                    rev_data.text_chars = len(rev.text)  

 
                     # generate revert data
                     revert = rev_detector.process(text_sha1, rev.id)
 
                     # generate revert data
                     revert = rev_detector.process(text_sha1, rev.id)


@@ -553,11 +581,6 @@ class WikiqParser():
                     rev_data.editor = rev.user.text 
                     rev_data.anon = rev.user.id is None
                 
                     rev_data.editor = rev.user.text 
                     rev_data.anon = rev.user.id is None
                 

-                #if re.match(r'^#redirect \[\[.*\]\]', rev.text, re.I):
-                #    redirect = True
-                #else:
-                #    redirect = False
-                

                 #TODO missing: additions_size deletions_size
                 
                 # if collapse user was on, lets run that
                 #TODO missing: additions_size deletions_size
                 
                 # if collapse user was on, lets run that


@@ -661,13 +684,14 @@ class WikiqParser():
 
         outtable = rows_to_table(self.parquet_buffer, self.schema)
         if self.pq_writer is None:
 
         outtable = rows_to_table(self.parquet_buffer, self.schema)
         if self.pq_writer is None:

-            self.pq_writer = pq.ParquetWriter(self.output_file, schema, flavor='spark')
+            self.pq_writer = pq.ParquetWriter(self.output_file, self.schema, flavor='spark')

 
         self.pq_writer.write_table(outtable)
         self.parquet_buffer = []
         
     # depending on if we are configured to write tsv or parquet, we'll call a different function.
     def print_rev_data(self, rev_data):
 
         self.pq_writer.write_table(outtable)
         self.parquet_buffer = []
         
     # depending on if we are configured to write tsv or parquet, we'll call a different function.
     def print_rev_data(self, rev_data):

+

         if self.output_parquet is False:
             printfunc = self.write_tsv_row
         else:
         if self.output_parquet is False:
             printfunc = self.write_tsv_row
         else:


@@ -683,6 +707,21 @@ class WikiqParser():
         line = rev_data.to_tsv_row()
         print(line, file=self.output_file)
 
         line = rev_data.to_tsv_row()
         print(line, file=self.output_file)
 

+def open_siteinfo(siteinfo_file):
+    if re.match(r'.*\.7z$', siteinfo_file):
+        cmd = ["7za", "x", "-so", siteinfo_file, "*.json"] 
+    elif re.match(r'.*\.gz$', siteinfo_file):
+        cmd = ["zcat", siteinfo_file] 
+    elif re.match(r'.*\.bz2$', siteinfo_file):
+        cmd = ["bzcat", "-dk", siteinfo_file] 
+
+    try:
+        input_file = Popen(cmd, stdout=PIPE).stdout
+    except NameError:
+        input_file = open(siteinfo_file, 'r')
+
+    return input_file
+

 
 def open_input_file(input_filename):
     if re.match(r'.*\.7z$', input_filename):
 
 def open_input_file(input_filename):
     if re.match(r'.*\.7z$', input_filename):


@@ -758,6 +797,11 @@ parser.add_argument('-CP', '--comment-pattern', dest="regex_match_comment", defa
 parser.add_argument('-CPl', '--comment-pattern-label', dest="regex_comment_label", default=None, type=str, action='append',
                     help="The label for the outputted column based on matching the regex in comments.")
 
 parser.add_argument('-CPl', '--comment-pattern-label', dest="regex_comment_label", default=None, type=str, action='append',
                     help="The label for the outputted column based on matching the regex in comments.")
 

+parser.add_argument('--SI', '--siteinfo', dest="siteinfo", default=None, type=str,
+                    help="Path to archive containing siteinfo json. This is required for resolving redirects")
+
+
+

 args = parser.parse_args()
 
 
 args = parser.parse_args()
 
 


@@ -800,6 +844,7 @@ if len(args.dumpfiles) > 0:
             filename = os.path.join(output_dir, os.path.basename(filename))
             output_file = get_output_filename(filename, parquet = output_parquet)
 
             filename = os.path.join(output_dir, os.path.basename(filename))
             output_file = get_output_filename(filename, parquet = output_parquet)
 

+        print(args.siteinfo, file=sys.stderr)

         wikiq = WikiqParser(input_file,
                             output_file,
                             collapse_user=args.collapse_user,
         wikiq = WikiqParser(input_file,
                             output_file,
                             collapse_user=args.collapse_user,


@@ -811,7 +856,8 @@ if len(args.dumpfiles) > 0:
                             regex_revision_label = args.regex_revision_label,
                             regex_match_comment = args.regex_match_comment,
                             regex_comment_label = args.regex_comment_label,
                             regex_revision_label = args.regex_revision_label,
                             regex_match_comment = args.regex_match_comment,
                             regex_comment_label = args.regex_comment_label,

-                            output_parquet=output_parquet)
+                            output_parquet=output_parquet,
+                            siteinfo_file = args.siteinfo)

 
         wikiq.process()
 
 
         wikiq.process()
 


@@ -830,7 +876,8 @@ else:
                         regex_match_revision = args.regex_match_revision,
                         regex_revision_label = args.regex_revision_label,
                         regex_match_comment = args.regex_match_comment,
                         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,
+                        siteinfo_file = args.siteinfo)

 
     wikiq.process() 
 
 
     wikiq.process()