Merge pull request allenai#139 from allenai/robust04-fix

.z compression support for robust04

Author: seanmacavaney

.github/workflows/push.yml

@@ -14,7 +14,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: [3.6, 3.7, 3.8, 3.9]
+        python-version: [3.7, 3.8, 3.9]
         os: ['ubuntu-latest', 'windows-latest', 'macOs-latest']
         architecture: ['x64']
 

ir_datasets/formats/trec.py

@@ -113,9 +113,15 @@ def docs_iter(self):
 
     def _docs_iter(self, path):
         if Path(path).is_file():
-            if str(path).endswith('.gz'):
+            path_suffix = Path(path).suffix.lower()
+            if path_suffix == '.gz':
                 with gzip.open(path, 'rb') as f:
                     yield from self._parser(f)
+            elif path_suffix in ['.z', '.0z', '.1z', '.2z']:
+                # unix "compress" command encoding
+                unlzw3 = ir_datasets.lazy_libs.unlzw3()
+                with io.BytesIO(unlzw3.unlzw(path)) as f:
+                    yield from self._parser(f)
             else:
                 with path.open('rb') as f:
                     yield from self._parser(f)

ir_datasets/lazy_libs.py

@@ -102,4 +102,10 @@ def pyautocorpus():
     if 'pyautocorpus' not in _cache:
         import pyautocorpus
         _cache['pyautocorpus'] = pyautocorpus
-    return _cache['pyautocorpus']    
+    return _cache['pyautocorpus']
+
+def unlzw3():
+    if 'unlzw3' not in _cache:
+        import unlzw3
+        _cache['unlzw3'] = unlzw3
+    return _cache['unlzw3']

requirements.txt

@@ -11,3 +11,4 @@ warc3-wet-clueweb09>=0.2.5
 zlib-state>=0.1.3
 ijson>=3.1.3
 pyautocorpus>=0.1.1
+unlzw3>=0.2.1