repace tar.bz2 by tar for SCSN dataset

yewsg · May 2, 2019 · 0270bb9 · 0270bb9
1 parent f5ba23a
commit 0270bb9
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 8 deletions.
diff --git a/yews/datasets/__init__.py b/yews/datasets/__init__.py
@@ -8,6 +8,7 @@
 from .files import FileDataset
 from .packaged_datasets import Mariana
 from .packaged_datasets import PackagedDataset
+from .packaged_datasets import SCSN
 from .packaged_datasets import Wenchuan
 from .sac import MarianaFromSource
 from .utils import *
@@ -23,5 +24,6 @@
     'PackagedDataset',
     'Wenchuan',
     'Mariana',
+    'SCSN',
     'MarianaFromSource',
 )
diff --git a/yews/datasets/packaged_datasets.py b/yews/datasets/packaged_datasets.py
@@ -1,9 +1,9 @@
 from .dirs import DatasetArrayFolder
-from .utils import extract_bz2
+from .utils import extract_tar
 from .utils import URL
 
 class PackagedDataset(DatasetArrayFolder):
-    """Packaged Dataset in .tar.bz2 foramt.
+    """Packaged Dataset in .tar.* foramt.
 
     Args:
         path (str): Root directory where ``dataset_name/samples.npy`
@@ -36,8 +36,8 @@ def __init__(self, download=False, **kwargs):
                 if not fpath.is_file():
                     url.download(self.root)
                 # extract file under root directory
-                print("Extracting dataset (this may minutes to hours) ...")
-                extract_bz2(fpath, self.root)
+                print("Extracting dataset (this may take minutes to hours) ...")
+                extract_tar(fpath, self.root)
                 # try initiate DatasetArrayFolder again
                 super().__init__(**kwargs)
             else:
@@ -99,4 +99,4 @@ class SCSN(PackagedDataset):
 
     """
 
-    url = 'https://www.dropbox.com/s/8r9rv9hudn7jt1q/scsn.tar.bz2?dl=1'
+    url = 'https://www.dropbox.com/s/u16lkjxliw8fouj/scsn.tar?dl=1'
diff --git a/yews/datasets/utils.py b/yews/datasets/utils.py
@@ -158,10 +158,11 @@ def download(self, root, filename=None):
 
 ################################################################################
 #
-#                           Utilities for Bz2
+#                           Utilities for Tarfile
 #
 ################################################################################
 
-def extract_bz2(infile, outdir='.'):
-    with tarfile.open(infile, 'r:bz2') as tar:
+
+def extract_tar(infile, outdir='.', mode='r:*'):
+    with tarfile.open(infile, mode=mode) as tar:
         tar.extractall(outdir)