wikiq

   1 #!/usr/bin/env python3
   2
   3 # original wikiq headers are: title articleid revid date_time anon
   4 # editor editor_id minor text_size text_entropy text_md5 reversion
   5 # additions_size deletions_size
   6
   7 import argparse
   8 import sys
   9 import os, os.path
  10 import re
  11 from datetime import datetime,timezone
  12
  13 from subprocess import Popen, PIPE
  14 from collections import deque
  15 from hashlib import sha1
  16
  17 from mwxml import Dump
  18
  19 from deltas.tokenizers import wikitext_split
  20 import mwpersistence
  21 import mwreverts
  22 from urllib.parse import quote
  23 TO_ENCODE = ('title', 'editor')
  24 PERSISTENCE_RADIUS=7
  25 from deltas import SequenceMatcher
  26 from deltas import SegmentMatcher
  27
  28 import dataclasses as dc
  29 from dataclasses import dataclass
  30 import pyarrow as pa
  31 import pyarrow.parquet as pq
  32
  33 class PersistMethod:
  34     none = 0
  35     sequence = 1
  36     segment = 2
  37     legacy = 3
  38
  39 def calculate_persistence(tokens_added):
  40     return(sum([(len(x.revisions)-1) for x in tokens_added]),
  41            len(tokens_added))
  42
  43 class WikiqIterator():
  44     def __init__(self, fh, collapse_user=False):
  45         self.fh = fh
  46         self.collapse_user = collapse_user
  47         self.mwiterator = Dump.from_file(self.fh)
  48         self.namespace_map = { ns.id : ns.name for ns in
  49                                self.mwiterator.site_info.namespaces }
  50         self.__pages = self.load_pages()
  51
  52     def load_pages(self):
  53         for page in self.mwiterator:
  54             yield WikiqPage(page,
  55                             namespace_map = self.namespace_map,
  56                             collapse_user=self.collapse_user)
  57
  58     def __iter__(self):
  59         return self.__pages
  60
  61     def __next__(self):
  62         return next(self._pages)
  63
  64 class WikiqPage():
  65     __slots__ = ('id', 'title', 'namespace', 'redirect',
  66                  'restrictions', 'mwpage', '__revisions',
  67                  'collapse_user')
  68
  69     def __init__(self, page, namespace_map, collapse_user=False):
  70         self.id = page.id
  71         self.namespace = page.namespace
  72         # following mwxml, we assume namespace 0 in cases where
  73         # page.namespace is inconsistent with namespace_map
  74         if page.namespace not in namespace_map:
  75             self.title = page.title
  76             page.namespace = 0
  77         if page.namespace != 0:
  78             self.title = ':'.join([namespace_map[page.namespace], page.title])
  79         else:
  80             self.title = page.title
  81         self.restrictions = page.restrictions
  82         self.collapse_user = collapse_user
  83         self.mwpage = page
  84         self.__revisions = self.rev_list()
  85
  86     def rev_list(self):
  87         # Outline for how we want to handle collapse_user=True
  88         # iteration   rev.user   prev_rev.user   add prev_rev?
  89         #         0          A            None           Never
  90         #         1          A               A           False
  91         #         2          B               A            True
  92         #         3          A               B            True
  93         #         4          A               A           False
  94         # Post-loop                          A          Always
  95         for i, rev in enumerate(self.mwpage):
  96             # never yield the first time
  97             if i == 0:
  98                 if self.collapse_user:
  99                     collapsed_revs = 1
 100                     rev.collapsed_revs = collapsed_revs
 101
 102             else:
 103                 if self.collapse_user:
 104                     # yield if this is the last edit in a seq by a user and reset
 105                     # also yield if we do know who the user is
 106
 107                     if rev.deleted.user or prev_rev.deleted.user:
 108                         yield prev_rev
 109                         collapsed_revs = 1
 110                         rev.collapsed_revs = collapsed_revs
 111
 112                     elif not rev.user.text == prev_rev.user.text:
 113                         yield prev_rev
 114                         collapsed_revs = 1
 115                         rev.collapsed_revs = collapsed_revs
 116                     # otherwise, add one to the counter
 117                     else:
 118                         collapsed_revs += 1
 119                         rev.collapsed_revs = collapsed_revs
 120                 # if collapse_user is false, we always yield
 121                 else:
 122                     yield prev_rev
 123
 124             prev_rev = rev
 125
 126         # also yield the final time
 127         yield prev_rev
 128
 129     def __iter__(self):
 130         return self.__revisions
 131
 132     def __next__(self):
 133         return next(self.__revisions)
 134
 135
 136 """
 137 A RegexPair is defined by a regular expression (pattern) and a label.
 138 The pattern can include capture groups.  If it does then each capture group will have a resulting column in the output.
 139 If the pattern does not include a capture group, then only one output column will result.
 140 """
 141 class RegexPair(object):
 142     def __init__(self, pattern, label):
 143         self.pattern = re.compile(pattern)
 144         self.label = label
 145         self.has_groups = bool(self.pattern.groupindex)
 146         if self.has_groups:
 147             self.capture_groups = list(self.pattern.groupindex.keys())
 148
 149     def get_pyarrow_fields(self):
 150         if self.has_groups:
 151             fields = [pa.field(self._make_key(cap_group),pa.list_(pa.string()))
 152                       for cap_group in self.capture_groups]
 153         else:
 154             fields = [pa.field(self.label, pa.list_(pa.string()))]
 155
 156         return fields
 157
 158     def _make_key(self, cap_group):
 159         return ("{}_{}".format(self.label, cap_group))
 160
 161     def matchmake(self, content, rev_data):
 162
 163         temp_dict = {}
 164         # if there are named capture groups in the regex
 165         if self.has_groups:
 166
 167             # if there are matches of some sort in this revision content, fill the lists for each cap_group
 168             if self.pattern.search(content) is not None:
 169                 m = self.pattern.finditer(content)
 170                 matchobjects = list(m)
 171
 172                 for cap_group in self.capture_groups:
 173                     key = self._make_key(cap_group)
 174                     temp_list = []
 175                     for match in matchobjects:
 176                         # we only want to add the match for the capture group if the match is not None
 177                         if match.group(cap_group) != None:
 178                             temp_list.append(match.group(cap_group))
 179
 180                     # if temp_list of matches is empty just make that column None
 181                     if len(temp_list)==0:
 182                         temp_dict[key] = None
 183                     # else we put in the list we made in the for-loop above
 184                     else:
 185                         temp_dict[key] = ', '.join(temp_list)
 186
 187             # there are no matches at all in this revision content, we default values to None
 188             else:
 189                 for cap_group in self.capture_groups:
 190                     key = self._make_key(cap_group)
 191                     temp_dict[key] = None
 192
 193         # there are no capture groups, we just search for all the matches of the regex
 194         else:
 195             #given that there are matches to be made
 196             if type(content) in(str, bytes):
 197                 if self.pattern.search(content) is not None:
 198                     m = self.pattern.findall(content)
 199                     temp_dict[self.label] = ', '.join(m)
 200                 else:
 201                     temp_dict[self.label] = None
 202
 203         # update rev_data with our new columns
 204         for k, v in temp_dict.items():
 205             setattr(rev_data, k, v)
 206
 207         return rev_data
 208
 209 """
 210
 211 We used to use a dictionary to collect fields for the output.
 212 Now we use dataclasses. Compared to a dictionary, this should help:
 213 - prevent some bugs
 214 - make it easier to output parquet data.
 215 - use class attribute '.' syntax instead of dictionary syntax.
 216 - improve support for tooling (autocomplete, type hints)
 217 - use type information to define formatting rules
 218
 219 Depending on the parameters passed into Wikiq, the output schema can be different.
 220 Therefore, we need to end up constructing a dataclass with the correct output schema.
 221 It also needs to have the correct pyarrow schema so we can write parquet files.
 222
 223 The RevDataBase type has all the fields that will be output no matter how wikiq is invoked.
 224 """
 225 @dataclass()
 226 class RevDataBase():
 227     revid: int
 228     date_time: datetime
 229     articleid: int
 230     editorid: int
 231     title: str
 232     namespace: int
 233     deleted: bool
 234     text_chars: int = None
 235     revert: bool = None
 236     reverteds: list[int] = None
 237     sha1: str = None
 238     minor: bool = None
 239     editor: str = None
 240     anon: bool = None
 241
 242     # toggles url encoding. this isn't a dataclass field since it doesn't have a type annotation
 243     urlencode = False
 244
 245     # defines pyarrow schema.
 246     # each field in the data class needs an entry in this array.
 247     # the names should match and be in the same order.
 248     # this isn't a dataclass field since it doesn't have a type annotation
 249     pa_schema_fields = [
 250         pa.field("revid", pa.int64()),
 251         pa.field("date_time", pa.timestamp('ms')),
 252         pa.field("articleid",pa.int64()),
 253         pa.field("editorid",pa.int64()),
 254         pa.field("title",pa.string()),
 255         pa.field("namespace",pa.int32()),
 256         pa.field("deleted",pa.bool_()),
 257         pa.field("test_chars",pa.int32()),
 258         pa.field("revert",pa.bool_()),
 259         pa.field("reverteds",pa.list_(pa.int64())),
 260         pa.field("sha1",pa.string()),
 261         pa.field("minor",pa.bool_()),
 262         pa.field("editor",pa.string()),
 263         pa.field("anon",pa.bool_())
 264     ]
 265
 266     # pyarrow is a columnar format, so most of the work happens in the flush_parquet_buffer function
 267     def to_pyarrow(self):
 268         return dc.astuple(self)
 269
 270     # logic to convert each field into the wikiq tsv format goes here.
 271     def to_tsv_row(self):
 272
 273         row = []
 274         for f in dc.fields(self):
 275             val = getattr(self, f.name)
 276             if getattr(self, f.name) is None:
 277                 row.append("")
 278             elif f.type == bool:
 279                 row.append("TRUE" if val else "FALSE")
 280
 281             elif f.type == datetime:
 282                 row.append(val.strftime('%Y-%m-%d %H:%M:%S'))
 283
 284             elif f.name in {'editor','title'}:
 285                 s = '"' + val + '"'
 286                 if self.urlencode and f.name in TO_ENCODE:
 287                     row.append(quote(str(s)))
 288                 else:
 289                     row.append(s)
 290
 291             elif f.type == list[int]:
 292                 row.append('"' + ",".join([str(x) for x in val]) + '"')
 293
 294             elif f.type == str:
 295                 if self.urlencode and f.name in TO_ENCODE:
 296                     row.append(quote(str(val)))
 297                 else:
 298                     row.append(val)
 299             else:
 300                 row.append(val)
 301
 302         return '\t'.join(map(str,row))
 303
 304     def header_row(self):
 305         return '\t'.join(map(lambda f: f.name, dc.fields(self)))
 306
 307 """
 308
 309 If collapse=True we'll use a RevDataCollapse dataclass.
 310 This class inherits from RevDataBase. This means that it has all the same fields and functions.
 311
 312 It just adds a new field and updates the pyarrow schema.
 313
 314 """
 315 @dataclass()
 316 class RevDataCollapse(RevDataBase):
 317     collapsed_revs:int = None
 318
 319     pa_collapsed_revs_schema = pa.field('collapsed_revs',pa.int64())
 320     pa_schema_fields = RevDataBase.pa_schema_fields + [pa_collapsed_revs_schema]
 321
 322 """
 323
 324 If persistence data is to be computed we'll need the fields added by RevDataPersistence.
 325
 326 """
 327 @dataclass()
 328 class RevDataPersistence(RevDataBase):
 329     token_revs:int = None
 330     tokens_added:int = None
 331     tokens_removed:int = None
 332     tokens_window:int = None
 333
 334     pa_persistence_schema_fields = [
 335         pa.field("token_revs", pa.int64()),
 336         pa.field("tokens_added", pa.int64()),
 337         pa.field("tokens_removed", pa.int64()),
 338         pa.field("tokens_window", pa.int64())]
 339
 340     pa_schema_fields = RevDataBase.pa_schema_fields  + pa_persistence_schema_fields
 341
 342 """
 343 class RevDataCollapsePersistence uses multiple inheritence to make a class that has both persistence and collapse fields.
 344
 345 """
 346 @dataclass()
 347 class RevDataCollapsePersistence(RevDataCollapse, RevDataPersistence):
 348     pa_schema_fields = RevDataCollapse.pa_schema_fields + RevDataPersistence.pa_persistence_schema_fields
 349
 350 class WikiqParser():
 351     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):
 352         """
 353         Parameters:
 354            persist : what persistence method to use. Takes a PersistMethod value
 355         """
 356         self.input_file = input_file
 357
 358         self.collapse_user = collapse_user
 359         self.persist = persist
 360         self.namespaces = []
 361         self.urlencode = urlencode
 362         self.revert_radius = revert_radius
 363
 364         if namespaces is not None:
 365             self.namespace_filter = set(namespaces)
 366         else:
 367             self.namespace_filter = None
 368
 369         self.regex_schemas = []
 370         self.regex_revision_pairs = self.make_matchmake_pairs(regex_match_revision, regex_revision_label)
 371         self.regex_comment_pairs = self.make_matchmake_pairs(regex_match_comment, regex_comment_label)
 372
 373
 374         # This is where we set the type for revdata.
 375
 376         if self.collapse_user is True:
 377             if self.persist == PersistMethod.none:
 378                 revdata_type = RevDataCollapse
 379             else:
 380                 revdata_type = RevDataCollapsePersistence
 381         elif self.persist != PersistMethod.none:
 382             revdata_type = RevDataPersistence
 383         else:
 384             revdata_type = RevDataBase
 385
 386         # if there are regex fields, we need to add them to the revdata type.
 387         regex_fields = [(field.name, list[str], dc.field(default=None)) for field in self.regex_schemas]
 388
 389         # make_dataclass is a function that defines a new dataclass type.
 390         # here we extend the type we have already chosen and add the regular expression types
 391         self.revdata_type = dc.make_dataclass('RevData_Parser',
 392                                               fields=regex_fields,
 393                                               bases=(revdata_type,))
 394
 395         # we also need to make sure that we have the right pyarrow schema
 396         self.revdata_type.pa_schema_fields = revdata_type.pa_schema_fields + self.regex_schemas
 397
 398         self.revdata_type.urlencode = self.urlencode
 399
 400         self.schema = pa.schema(self.revdata_type.pa_schema_fields)
 401
 402         # here we initialize the variables we need for output.
 403         if output_parquet is True:
 404             self.output_parquet = True
 405             self.pq_writer = None
 406             self.output_file = output_file
 407             self.parquet_buffer = []
 408             self.parquet_buffer_size = parquet_buffer_size
 409         else:
 410             self.print_header = True
 411             if output_file == sys.stdout:
 412
 413                 self.output_file = output_file
 414             else:
 415                 self.output_file = open(output_file,'w')
 416             self.output_parquet = False
 417
 418     def make_matchmake_pairs(self, patterns, labels):
 419         if (patterns is not None and labels is not None) and \
 420            (len(patterns) == len(labels)):
 421             result = []
 422             for pattern, label in zip(patterns, labels):
 423                 rp = RegexPair(pattern, label)
 424                 result.append(rp)
 425                 self.regex_schemas = self.regex_schemas + rp.get_pyarrow_fields()
 426             return result
 427         elif (patterns is None and labels is None):
 428             return []
 429         else:
 430             sys.exit('Each regular expression *must* come with a corresponding label and vice versa.')
 431
 432     def matchmake_revision(self, rev, rev_data):
 433         rev_data = self.matchmake_text(rev.text, rev_data)
 434         rev_data = self.matchmake_comment(rev.comment, rev_data)
 435         return rev_data
 436
 437     def matchmake_text(self, text, rev_data):
 438          return self.matchmake_pairs(text, rev_data, self.regex_revision_pairs)
 439
 440     def matchmake_comment(self, comment, rev_data):
 441         return self.matchmake_pairs(comment, rev_data, self.regex_comment_pairs)
 442
 443     def matchmake_pairs(self, text, rev_data, pairs):
 444         for pair in pairs:
 445             rev_data = pair.matchmake(text, rev_data)
 446         return rev_data
 447
 448     def __get_namespace_from_title(self, title):
 449         default_ns = None
 450
 451         for ns in self.namespaces:
 452             # skip if the namespace is not defined
 453             if ns == None:
 454                 default_ns = self.namespaces[ns]
 455                 continue
 456
 457             if title.startswith(ns + ":"):
 458                 return self.namespaces[ns]
 459
 460         # if we've made it this far with no matches, we return the default namespace
 461         return default_ns
 462
 463
 464     def process(self):
 465
 466         # create a regex that creates the output filename
 467         # output_filename = re.sub(r'^.*/(enwiki\-\d+)\-.*p(\d+)p.*$',
 468         #                         r'output/wikiq-\1-\2.tsv',
 469         #                         input_filename)
 470
 471         # Construct dump file iterator
 472         dump = WikiqIterator(self.input_file, collapse_user=self.collapse_user)
 473
 474         # extract list of namspaces
 475         self.namespaces = {ns.name : ns.id for ns in dump.mwiterator.site_info.namespaces}
 476
 477         page_count = 0
 478         rev_count = 0
 479
 480
 481         # Iterate through pages
 482         for page in dump:
 483             namespace = page.namespace if page.namespace is not None else self.__get_namespace_from_title(page.title)
 484
 485             # skip namespaces not in the filter
 486             if self.namespace_filter is not None:
 487                 if namespace not in self.namespace_filter:
 488                     continue
 489
 490             rev_detector = mwreverts.Detector(radius = self.revert_radius)
 491
 492             if self.persist != PersistMethod.none:
 493                 window = deque(maxlen=PERSISTENCE_RADIUS)
 494
 495                 if self.persist == PersistMethod.sequence:
 496                     state = mwpersistence.DiffState(SequenceMatcher(tokenizer = wikitext_split),
 497                                                     revert_radius=PERSISTENCE_RADIUS)
 498
 499                 elif self.persist == PersistMethod.segment:
 500                     state = mwpersistence.DiffState(SegmentMatcher(tokenizer = wikitext_split),
 501                                                     revert_radius=PERSISTENCE_RADIUS)
 502
 503                 # self.persist == PersistMethod.legacy
 504                 else:
 505                     from mw.lib import persistence
 506                     state = persistence.State()
 507
 508             # Iterate through a page's revisions
 509             for rev in page:
 510
 511                 # create a new data object instead of a dictionary.
 512                 rev_data = self.revdata_type(revid = rev.id,
 513                                              date_time = datetime.fromtimestamp(rev.timestamp.unix(), tz=timezone.utc),
 514                                              articleid = page.id,
 515                                              editorid = "" if rev.deleted.user == True or rev.user.id is None else rev.user.id,
 516                                              title =  page.title,
 517                                              deleted = rev.deleted.text,
 518                                              namespace = namespace
 519                                              )
 520
 521                 rev_data = self.matchmake_revision(rev, rev_data)
 522
 523                 if not rev.deleted.text:
 524                     # rev.text can be None if the page has no text
 525                     if not rev.text:
 526                         rev.text = ""
 527                     # if text exists, we'll check for a sha1 and generate one otherwise
 528
 529                     if rev.sha1:
 530                         text_sha1 = rev.sha1
 531                     else:
 532                         text_sha1 = sha1(bytes(rev.text, "utf8")).hexdigest()
 533
 534                     rev_data.sha1 = text_sha1
 535
 536                     # TODO rev.bytes doesn't work.. looks like a bug
 537                     rev_data.text_chars = len(rev.text)
 538
 539                     # generate revert data
 540                     revert = rev_detector.process(text_sha1, rev.id)
 541
 542                     if revert:
 543                         rev_data.revert = True
 544                         rev_data.reverteds = revert.reverteds
 545                     else:
 546                         rev_data.revert = False
 547
 548                 # if the fact that the edit was minor can be hidden, this might be an issue
 549                 rev_data.minor = rev.minor
 550
 551                 if not rev.deleted.user:
 552                     # wrap user-defined editors in quotes for fread
 553                     rev_data.editor = rev.user.text
 554                     rev_data.anon = rev.user.id is None
 555
 556                 #if re.match(r'^#redirect \[\[.*\]\]', rev.text, re.I):
 557                 #    redirect = True
 558                 #else:
 559                 #    redirect = False
 560
 561                 #TODO missing: additions_size deletions_size
 562
 563                 # if collapse user was on, lets run that
 564                 if self.collapse_user:
 565                     rev_data.collapsed_revs = rev.collapsed_revs
 566
 567                 # get the
 568                 if self.persist != PersistMethod.none:
 569                     if not rev.deleted.text:
 570
 571                         if self.persist != PersistMethod.legacy:
 572                             _, tokens_added, tokens_removed = state.update(rev.text, rev.id)
 573
 574                         else:
 575                             _, tokens_added, tokens_removed = state.process(rev.text, rev.id, text_sha1)
 576
 577                         window.append((rev.id, rev_data, tokens_added, tokens_removed))
 578
 579                         if len(window) == PERSISTENCE_RADIUS:
 580                             old_rev_id, old_rev_data, old_tokens_added, old_tokens_removed = window[0]
 581
 582                             num_token_revs, num_tokens = calculate_persistence(old_tokens_added)
 583
 584                             old_rev_data.token_revs = num_token_revs
 585                             old_rev_data.tokens_added = num_tokens
 586                             old_rev_data.tokens_removed = len(old_tokens_removed)
 587                             old_rev_data.tokens_window = PERSISTENCE_RADIUS-1
 588
 589                             self.print_rev_data(old_rev_data)
 590
 591                 else:
 592                     self.print_rev_data(rev_data)
 593
 594                 rev_count += 1
 595
 596             if self.persist != PersistMethod.none:
 597                 # print out metadata for the last RADIUS revisions
 598                 for i, item in enumerate(window):
 599                     # if the window was full, we've already printed item 0
 600                     if len(window) == PERSISTENCE_RADIUS and i == 0:
 601                         continue
 602
 603                     rev_id, rev_data, tokens_added, tokens_removed = item
 604                     num_token_revs, num_tokens = calculate_persistence(tokens_added)
 605
 606                     rev_data.token_revs = num_token_revs
 607                     rev_data.tokens_added = num_tokens
 608                     rev_data.tokens_removed = len(tokens_removed)
 609                     rev_data.tokens_window = len(window)-(i+1)
 610                     self.print_rev_data(rev_data)
 611
 612             page_count += 1
 613
 614         print("Done: %s revisions and %s pages." % (rev_count, page_count),
 615               file=sys.stderr)
 616
 617         # remember to flush the parquet_buffer if we're done
 618         if self.output_parquet is True:
 619             self.flush_parquet_buffer()
 620             self.pq_writer.close()
 621
 622         else:
 623             self.output_file.close()
 624
 625
 626     """
 627     For performance reasons it's better to write parquet in batches instead of one row at a time.
 628     So this function just puts the data on a buffer. If the buffer is full, then it gets flushed (written).
 629     """
 630     def write_parquet_row(self, rev_data):
 631         padata = rev_data.to_pyarrow()
 632         self.parquet_buffer.append(padata)
 633
 634         if len(self.parquet_buffer) >= self.parquet_buffer_size:
 635             self.flush_parquet_buffer()
 636
 637
 638     """
 639     Function that actually writes data to the parquet file.
 640     It needs to transpose the data from row-by-row to column-by-column
 641     """
 642     def flush_parquet_buffer(self):
 643
 644         """
 645         Returns the pyarrow table that we'll write
 646         """
 647         def rows_to_table(rg, schema):
 648             cols = []
 649             first = rg[0]
 650             for col in first:
 651                 cols.append([col])
 652
 653             for row in rg[1:]:
 654                 for j in range(len(cols)):
 655                     cols[j].append(row[j])
 656
 657             arrays = []
 658             for col, typ in zip(cols, schema.types):
 659                 arrays.append(pa.array(col, typ))
 660             return pa.Table.from_arrays(arrays, schema=schema)
 661
 662         outtable = rows_to_table(self.parquet_buffer, self.schema)
 663         if self.pq_writer is None:
 664             self.pq_writer = pq.ParquetWriter(self.output_file, schema, flavor='spark')
 665
 666         self.pq_writer.write_table(outtable)
 667         self.parquet_buffer = []
 668
 669     # depending on if we are configured to write tsv or parquet, we'll call a different function.
 670     def print_rev_data(self, rev_data):
 671         if self.output_parquet is False:
 672             printfunc = self.write_tsv_row
 673         else:
 674             printfunc = self.write_parquet_row
 675
 676         printfunc(rev_data)
 677
 678     def write_tsv_row(self, rev_data):
 679         if self.print_header:
 680             print(rev_data.header_row(), file=self.output_file)
 681             self.print_header = False
 682
 683         line = rev_data.to_tsv_row()
 684         print(line, file=self.output_file)
 685
 686
 687 def open_input_file(input_filename):
 688     if re.match(r'.*\.7z$', input_filename):
 689         cmd = ["7za", "x", "-so", input_filename, "*.xml"]
 690     elif re.match(r'.*\.gz$', input_filename):
 691         cmd = ["zcat", input_filename]
 692     elif re.match(r'.*\.bz2$', input_filename):
 693         cmd = ["bzcat", "-dk", input_filename]
 694
 695     try:
 696         input_file = Popen(cmd, stdout=PIPE).stdout
 697     except NameError:
 698         input_file = open(input_filename, 'r')
 699
 700     return input_file
 701
 702 def get_output_filename(input_filename, parquet = False):
 703     output_filename = re.sub(r'\.(7z|gz|bz2)?$', '', input_filename)
 704     output_filename = re.sub(r'\.xml', '', output_filename)
 705     if parquet is False:
 706         output_filename = output_filename + ".tsv"
 707     else:
 708         output_filename = output_filename + ".parquet"
 709     return output_filename
 710
 711 def open_output_file(input_filename):
 712     # create a regex that creates the output filename
 713     output_filename = get_output_filename(input_filename, parquet = False)
 714     output_file = open(output_filename, "w")
 715     return output_file
 716
 717 parser = argparse.ArgumentParser(description='Parse MediaWiki XML database dumps into tab delimitted data.')
 718
 719 # arguments for the input direction
 720 parser.add_argument('dumpfiles', metavar="DUMPFILE", nargs="*", type=str,
 721                     help="Filename of the compressed or uncompressed XML database dump. If absent, we'll look for content on stdin and output on stdout.")
 722
 723 parser.add_argument('-o', '--output-dir', metavar='DIR', dest='output_dir', type=str, nargs=1,
 724                     help="Directory for output files. If it ends with .parquet output will be in parquet format.")
 725
 726 parser.add_argument('-s', '--stdout', dest="stdout", action="store_true",
 727                     help="Write output to standard out (do not create dump file)")
 728
 729 parser.add_argument('--collapse-user', dest="collapse_user", action="store_true",
 730                     help="Operate only on the final revision made by user a user within all sequences of consecutive edits made by a user. This can be useful for addressing issues with text persistence measures.")
 731
 732 parser.add_argument('-p', '--persistence', dest="persist", default=None, const='', type=str, choices = ['','segment','sequence','legacy'], nargs='?',
 733                     help="Compute and report measures of content persistent: (1) persistent token revisions, (2) tokens added, and (3) number of revision used in computing the first measure. This may by slow.  The defualt is -p=sequence, which uses the same algorithm as in the past, but with improvements to wikitext parsing. Use -p=legacy for old behavior used in older research projects. Use -p=segment for advanced persistence calculation method that is robust to content moves, but prone to bugs, and slower.")
 734
 735 parser.add_argument('-u', '--url-encode', dest="urlencode", action="store_true",
 736                     help="Output url encoded text strings. This works around some data issues like newlines in editor names. In the future it may be used to output other text data.")
 737
 738 parser.add_argument('-n', '--namespace-include', dest="namespace_filter", type=int, action='append',
 739                     help="Id number of namspace to include. Can be specified more than once.")
 740
 741 parser.add_argument('-rr',
 742                     '--revert-radius',
 743                     dest="revert_radius",
 744                     type=int,
 745                     action='store',
 746                     default=15,
 747                     help="Number of edits to check when looking for reverts (default: 15)")
 748
 749 parser.add_argument('-RP', '--revision-pattern', dest="regex_match_revision", default=None, type=str, action='append',
 750                     help="The regular expression to search for in revision text. The regex must be surrounded by quotes.")
 751
 752 parser.add_argument('-RPl', '--revision-pattern-label', dest="regex_revision_label", default=None, type=str, action='append',
 753                     help="The label for the outputted column based on matching the regex in revision text.")
 754
 755 parser.add_argument('-CP', '--comment-pattern', dest="regex_match_comment", default=None, type=str, action='append',
 756                     help="The regular expression to search for in comments of revisions.")
 757
 758 parser.add_argument('-CPl', '--comment-pattern-label', dest="regex_comment_label", default=None, type=str, action='append',
 759                     help="The label for the outputted column based on matching the regex in comments.")
 760
 761 args = parser.parse_args()
 762
 763
 764
 765 # set persistence method
 766
 767 if args.persist is None:
 768     persist = PersistMethod.none
 769 elif args.persist == "segment":
 770     persist = PersistMethod.segment
 771 elif args.persist == "legacy":
 772     persist = PersistMethod.legacy
 773 else:
 774     persist = PersistMethod.sequence
 775
 776 if args.namespace_filter is not None:
 777     namespaces = args.namespace_filter
 778 else:
 779     namespaces = None
 780
 781 if len(args.dumpfiles) > 0:
 782     output_parquet = False
 783     for filename in args.dumpfiles:
 784         input_file = open_input_file(filename)
 785
 786         # open directory for output
 787         if args.output_dir:
 788             output_dir = args.output_dir[0]
 789         else:
 790             output_dir = "."
 791
 792         if output_dir.endswith(".parquet"):
 793             output_parquet = True
 794
 795         print("Processing file: %s" % filename, file=sys.stderr)
 796
 797         if args.stdout:
 798             output_file = sys.stdout
 799         else:
 800             filename = os.path.join(output_dir, os.path.basename(filename))
 801             output_file = get_output_filename(filename, parquet = output_parquet)
 802
 803         wikiq = WikiqParser(input_file,
 804                             output_file,
 805                             collapse_user=args.collapse_user,
 806                             persist=persist,
 807                             urlencode=args.urlencode,
 808                             namespaces=namespaces,
 809                             revert_radius=args.revert_radius,
 810                             regex_match_revision = args.regex_match_revision,
 811                             regex_revision_label = args.regex_revision_label,
 812                             regex_match_comment = args.regex_match_comment,
 813                             regex_comment_label = args.regex_comment_label,
 814                             output_parquet=output_parquet)
 815
 816         wikiq.process()
 817
 818         # close things
 819         input_file.close()
 820
 821 else:
 822     wikiq = WikiqParser(sys.stdin,
 823                         sys.stdout,
 824                         collapse_user=args.collapse_user,
 825                         persist=persist,
 826                         #persist_legacy=args.persist_legacy,
 827                         urlencode=args.urlencode,
 828                         namespaces=namespaces,
 829                         revert_radius=args.revert_radius,
 830                         regex_match_revision = args.regex_match_revision,
 831                         regex_revision_label = args.regex_revision_label,
 832                         regex_match_comment = args.regex_match_comment,
 833                         regex_comment_label = args.regex_comment_label)
 834
 835     wikiq.process()
 836
 837 # stop_words = "a,able,about,across,after,all,almost,also,am,among,an,and,any,are,as,at,be,because,been,but,by,can,cannot,could,dear,did,do,does,either,else,ever,every,for,from,get,got,had,has,have,he,her,hers,him,his,how,however,i,if,in,into,is,it,its,just,least,let,like,likely,may,me,might,most,must,my,neither,no,nor,not,of,off,often,on,only,or,other,our,own,rather,said,say,says,she,should,since,so,some,than,that,the,their,them,then,there,these,they,this,tis,to,too,twas,us,wants,was,we,were,what,when,where,which,while,who,whom,why,will,with,would,yet,you,your"
 838 # stop_words = stop_words.split(",")