]> code.communitydata.science - cdsc_reddit.git/blobdiff - helper.py
Add code for running tf-idf at the weekly level.
[cdsc_reddit.git] / helper.py
index 4dc6210bef50e8142f2c56fbc788700cc761a082..8f1dfe2e79bdd3ad29018e34adde7a30e2cac823 100644 (file)
--- a/helper.py
+++ b/helper.py
@@ -14,19 +14,11 @@ def find_dumps(dumpdir, base_pattern):
         fname, ext = path.splitext(fpath)
         dumpext[fname].append(ext)
 
-    ext_priority = ['.zst','.xz','.bz2']
+    ext_priority = ['.zst','.xz','.bz2','.gz']
 
     for base, exts in dumpext.items():
-        found = False
-        if len(exts) == 1:
-            yield base + exts[0]
-            found = True
-        else:
-            for ext in ext_priority:
-                if ext in exts:
-                    yield base + ext
-                    found = True
-        assert(found == True)
+        ext = [ext for ext in ext_priority if ext in exts][0]
+        yield base + ext
 
 def open_fileset(files):
     for fh in files:
@@ -48,6 +40,8 @@ def open_input_file(input_filename):
         cmd = ["xzcat",'-dk', '-T 20',input_filename]
     elif re.match(r'.*\.zst',input_filename):
         cmd = ['zstd','-dck', input_filename]
+    elif re.match(r'.*\.gz',input_filename):
+        cmd = ['gzip','-dc', input_filename]
     try:
         input_file = Popen(cmd, stdout=PIPE).stdout
     except NameError as e:

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