]> code.communitydata.science - mediawiki_dump_tools.git/blobdiff - wikiq
refactor regex matching in a tidier object oriented style
[mediawiki_dump_tools.git] / wikiq
diff --git a/wikiq b/wikiq
index 6a1ada5cccee93edc9ae4c96f374542b1b4ba118..7e75eda22228cd5c44c248c264c8435b3268a948 100755 (executable)
--- a/wikiq
+++ b/wikiq
@@ -34,55 +34,6 @@ def calculate_persistence(tokens_added):
     return(sum([(len(x.revisions)-1) for x in tokens_added]),
            len(tokens_added))
 
-def matchmake(scanned_content, rev_data, regex, label):
-    p = re.compile(regex)
-
-    temp_dict = {}
-    # if there are named capture groups in the regex
-    if bool(p.groupindex):
-        capture_groups = list(p.groupindex.keys())
-
-        # initialize the {capture_group_name:list} for each capture group
-        for cap_group in capture_groups:
-            temp_dict["{}_{}".format(label, cap_group)] = []
-
-        # if there are matches of some sort in this revision content, fill the lists for each cap_group
-        if p.search(scanned_content) is not None:
-            m = re.finditer(p,scanned_content)
-            matchobjects = list(m)
-
-            for cap_group in capture_groups:
-                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) != None:
-                        temp_list.append(match.group(cap_group))
-
-                # if temp_list of matches is empty just make that column None
-                if len(temp_list)==0:
-                    temp_dict["{}_{}".format(label, cap_group)] = None
-                # else we put in the list we made in the for-loop above
-                else:
-                    temp_dict["{}_{}".format(label, cap_group)] = ', '.join(temp_list)
-        
-        # there are no matches at all in this revision content, we default values to None
-        else:
-            for cap_group in capture_groups:
-                temp_dict["{}_{}".format(label, cap_group)] = 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 p.search(scanned_content) is not None:
-            m = p.findall(scanned_content)
-            temp_dict[label] = ', '.join(m)
-        else:
-            temp_dict[label] = None    
-    # update rev_data with our new columns
-    rev_data.update(temp_dict)
-#    print(rev_data.keys())
-    return rev_data
-
 
 class WikiqIterator():
     def __init__(self, fh, collapse_user=False):
@@ -176,13 +127,70 @@ class WikiqPage():
     def __next__(self):
         return next(self.__revisions)
 
+
+class RegexPair(object):
+    def __init__(self, pattern, label):
+        self.pattern = re.compile(pattern)
+        self.label = label
+        self.has_groups = bool(self.pattern.groupindex)
+        if self.has_groups:
+            self.capture_groups = list(self.pattern.groupindex.keys())
+            
+    def _make_key(self, cap_group):
+        return ("{}_{}".format(self.label, cap_group))
+
+    def matchmake(self, content, rev_data):
+        
+        temp_dict = {}
+        # if there are named capture groups in the regex
+        if self.has_groups:
+            # initialize the {capture_group_name:list} for each capture group
+            # if there are matches of some sort in this revision content, fill the lists for each cap_group
+            if self.pattern.search(content) is not None:
+                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) != None:
+                            temp_list.append(match.group(cap_group))
+
+                    # if temp_list of matches is empty just make that column None
+                    if len(temp_list)==0:
+                        temp_dict[key] = None
+                    # else we put in the list we made in the for-loop above
+                    else:
+                        temp_dict[key] = ', '.join(temp_list)
+
+            # there are no matches at all in this revision content, we default values to None
+            else:
+                for cap_group in self.capture_groups:
+                    key = self._make_key(cap_group)
+                    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 self.pattern.search(content) is not None:
+                m = self.pattern.findall(content)
+                temp_dict[self.label] = ', '.join(m)
+            else:
+                temp_dict[self.label] = None    
+        # update rev_data with our new columns
+        rev_data.update(temp_dict)
+        return rev_data
+
+
+        
 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):
         """ 
         Parameters:
            persist : what persistence method to use. Takes a PersistMethod value
         """
-
         self.input_file = input_file
         self.output_file = output_file
         self.collapse_user = collapse_user
@@ -191,16 +199,41 @@ class WikiqParser():
         self.namespaces = []
         self.urlencode = urlencode
         self.revert_radius = revert_radius
-        self.regex_match_revision = regex_match_revision
-        self.regex_revision_label = regex_revision_label
-        self.regex_match_comment = regex_match_comment
-        self.regex_comment_label = regex_comment_label
 
         if namespaces is not None:
             self.namespace_filter = set(namespaces)
         else:
             self.namespace_filter = None
 
+        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)
+        
+
+    def make_matchmake_pairs(self, patterns, 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)]
+        elif (patterns is None and labels is None):
+            return []
+        else:
+            sys.exit('Each regular expression *must* come with a corresponding label and vice versa.')
+
+    def matchmake(self, rev, rev_data):
+        rev_data = self.matchmake_revision(rev.text, rev_data)
+        rev_data = self.matchmake_comment(rev.comment, rev_data)
+        return rev_data
+
+    def matchmake_revision(self, text, rev_data):
+        return self.matchmake_pairs(text, rev_data, self.regex_revision_pairs)
+
+    def matchmake_comment(self, comment, rev_data):
+        return self.matchmake_pairs(comment, rev_data, self.regex_comment_pairs)
+
+    def matchmake_pairs(self, text, rev_data, pairs):
+        for pair in pairs:
+            rev_data = pair.matchmake(text, rev_data)
+        return rev_data
+
     def __get_namespace_from_title(self, title):
         default_ns = None
 
@@ -243,10 +276,8 @@ class WikiqParser():
                 if namespace not in self.namespace_filter:
                     continue
 
-#            print(self.revert_radius)
             rev_detector = mwreverts.Detector(radius = self.revert_radius)
 
-
             if self.persist != PersistMethod.none:
                 window = deque(maxlen=PERSISTENCE_RADIUS)
 
@@ -267,50 +298,17 @@ class WikiqParser():
             for rev in page:
                 
                 # initialize rev_data
-                rev_data = {}
-
-                # if the command line args only gave a label (and no regular expression is given)
-                if (self.regex_revision_label != None and self.regex_match_revision == None) or (self.regex_comment_label != None and self.regex_match_comment == None):
-                    sys.exit('The given regex label(s) has no corresponding regex to search for.')
-                
-                # if there's anything in the list of regex_match_revision
-                if self.regex_match_revision is not None:
-                    if (self.regex_revision_label == None) or (len(self.regex_match_revision) != len(self.regex_revision_label)):
-                        sys.exit('Each regular expression *must* come with a corresponding label and vice versa.')
-                    
-                    # initialize and construct the list of regex-label tuples
-                    pairs = []
-                    for i in range(0,len(self.regex_match_revision)):
-                        pairs.append((self.regex_match_revision[i], self.regex_revision_label[i]))
-
-                    # for each regex/label pair, we now run matchmake to check and output columns
-                    for pair in pairs:
-                        # pair[0] corresponds to the regex, pair[1] to the label
-                        rev_data = matchmake(rev.text, rev_data, pair[0], pair[1])
-                
-                # if there's anything in the list of regex_match_comment
-                if self.regex_match_comment is not None:
-                    if (self.regex_comment_label == None) or (len(self.regex_match_comment) != len(self.regex_comment_label)):
-                        sys.exit('Each regular expression *must* come with a corresponding label and vice versa.')
-                    
-                    # initialize and construct the list of regex-label tuples
-                    pairs = []
-                    for i in range(0,len(self.regex_match_comment)):
-                        pairs.append((self.regex_match_comment[i], self.regex_comment_label[i]))
-
-                    # for each regex/label pair, we now run matchmake to check and output columns
-                    for pair in pairs:
-                        # pair[0] corresponds to the regex, pair[1] to the label
-                        rev_data = matchmake(rev.comment, rev_data, pair[0], pair[1])
-
-                # we fill out the rest of the data structure now
-                rev_data['revid'] = rev.id
-                rev_data['date_time'] = rev.timestamp.strftime('%Y-%m-%d %H:%M:%S')
-                rev_data['articleid'] = page.id
-                rev_data['editor_id'] = "" if rev.deleted.user == True or rev.user.id is None else rev.user.id
-                rev_data['title'] = '"' + page.title + '"'
-                rev_data['namespace'] = namespace
-                rev_data['deleted'] = "TRUE" if rev.deleted.text else "FALSE"
+                rev_data = {
+                    'revid':rev.id,
+                    'date_time' : rev.timestamp.strftime('%Y-%m-%d %H:%M:%S'),
+                    'articleid' : page.id,
+                    'editor_id' : "" if rev.deleted.user == True or rev.user.id is None else rev.user.id,
+                    'title' : '"' + page.title + '"',
+                    'namespace' : namespace,
+                    'deleted' : "TRUE" if rev.deleted.text else "FALSE"
+                }
+
+                rev_data = self.matchmake(rev, rev_data)
 
                 # if revisions are deleted, /many/ things will be missing
                 if rev.deleted.text:

Community Data Science Collective || Want to submit a patch?