Repository: fatiando/pooch

Issue:
Add an Untar processor
**Description of the desired feature**

The `pooch.Unzip` processor unpacks a zip archive to a folder in the local file system after it's downloaded. We need a `pooch.Untar` processor that does the same thing but for `.tar` archives. The [tarfile](https://docs.python.org/3/library/tarfile.html) module can be used with the bonus that it supports compressed archives out of the box.

Comments:

Patch:
diff --git a/doc/api/index.rst b/doc/api/index.rst
index b4d0521..0c5a4d2 100644
--- a/doc/api/index.rst
+++ b/doc/api/index.rst
@@ -42,6 +42,7 @@ Processors
    :toctree: generated/
 
    Unzip
+   Untar
 
 Miscellaneous
 -------------
diff --git a/pooch/__init__.py b/pooch/__init__.py
index d4160ed..514b635 100644
--- a/pooch/__init__.py
+++ b/pooch/__init__.py
@@ -4,7 +4,7 @@ from . import version
 from .core import Pooch, create
 from .utils import os_cache, file_hash, make_registry, check_version
 from .downloaders import HTTPDownloader
-from .processors import Unzip
+from .processors import Unzip, Untar
 
 
 def test(doctest=True, verbose=True, coverage=False):
diff --git a/pooch/processors.py b/pooch/processors.py
index 3092d9e..efa2a1a 100644
--- a/pooch/processors.py
+++ b/pooch/processors.py
@@ -5,33 +5,35 @@ Post-processing hooks for Pooch.fetch
 import os
 
 from zipfile import ZipFile
+from tarfile import TarFile
 from warnings import warn
 
 
-class Unzip:  # pylint: disable=too-few-public-methods
+class ExtractorProcessor:  # pylint: disable=too-few-public-methods
     """
-    Processor that unpacks a zip archive and returns a list of all files.
+    Base class for extractions from compressed archives.
 
-    Use with :meth:`pooch.Pooch.fetch` to unzip a downloaded data file into a folder in
-    the local data store. :meth:`~pooch.Pooch.fetch` will return a list with the names
-    of the unzipped files instead of the zip archive.
+    Subclasses can be used with :meth:`pooch.Pooch.fetch` to unzip a downloaded
+    data file into a folder in the local data store. :meth:`~pooch.Pooch.fetch`
+    will return a list with the names of the extracted files instead of the
+    archive.
 
     Parameters
     ----------
     members : list or None
-        If None, will unpack all files in the zip archive. Otherwise, *members* must be
-        a list of file names to unpack from the archive. Only these files will be
-        unpacked.
+        If None, will unpack all files in the archive. Otherwise, *members*
+        must be a list of file names to unpack from the archive. Only these
+        files will be unpacked.
     """
 
+    suffix = None  # String appended to unpacked archive. To be implemented in subclass
+
     def __init__(self, members=None):
         self.members = members
 
     def __call__(self, fname, action, pooch):
         """
-        Extract all files from the given zip archive.
-
-        The output folder is ``{fname}.unzip``.
+        Extract all files from the given archive.
 
         Parameters
         ----------
@@ -50,36 +52,125 @@ class Unzip:  # pylint: disable=too-few-public-methods
         Returns
         -------
         fnames : list of str
-            A list of the full path to all files in the unzipped archive.
+            A list of the full path to all files in the extracted archive.
 
         """
-        unzipped = fname + ".unzip"
-        if action in ("update", "download") or not os.path.exists(unzipped):
-            # Make sure that the folder with the unzipped files exists
-            if not os.path.exists(unzipped):
-                os.makedirs(unzipped)
-            with ZipFile(fname, "r") as zip_file:
-                if self.members is None:
-                    warn("Unzipping contents of '{}' to '{}'".format(fname, unzipped))
-                    # Unpack all files from the archive into our new folder
-                    zip_file.extractall(path=unzipped)
-                else:
-                    for member in self.members:
-                        warn(
-                            "Unzipping '{}' from '{}' to '{}'".format(
-                                member, fname, unzipped
-                            )
-                        )
-                        # Extract the data file from within the archive
-                        with zip_file.open(member) as data_file:
-                            # Save it to our desired file name
-                            with open(os.path.join(unzipped, member), "wb") as output:
-                                output.write(data_file.read())
+        if self.suffix is None:
+            raise NotImplementedError(
+                "Derived classes must define the 'suffix' attribute."
+            )
+        extract_dir = fname + self.suffix
+        if action in ("update", "download") or not os.path.exists(extract_dir):
+            # Make sure that the folder with the extracted files exists
+            if not os.path.exists(extract_dir):
+                os.makedirs(extract_dir)
+            self._extract_file(fname, extract_dir)
         # Get a list of all file names (including subdirectories) in our folder of
         # unzipped files.
         fnames = [
             os.path.join(path, fname)
-            for path, _, files in os.walk(unzipped)
+            for path, _, files in os.walk(extract_dir)
             for fname in files
         ]
         return fnames
+
+    def _extract_file(self, fname, extract_dir):
+        """
+        This method receives an argument for the archive to extract and the
+        destination path. MUST BE IMPLEMENTED BY CHILD CLASSES.
+        """
+        raise NotImplementedError
+
+
+class Unzip(ExtractorProcessor):  # pylint: disable=too-few-public-methods
+    """
+    Processor that unpacks a zip archive and returns a list of all files.
+
+    Use with :meth:`pooch.Pooch.fetch` to unzip a downloaded data file into a folder in
+    the local data store. :meth:`~pooch.Pooch.fetch` will return a list with the names
+    of the unzipped files instead of the zip archive.
+
+    The output folder is ``{fname}.unzip``.
+
+
+    Parameters
+    ----------
+    members : list or None
+        If None, will unpack all files in the zip archive. Otherwise, *members* must be
+        a list of file names to unpack from the archive. Only these files will be
+        unpacked.
+    """
+
+    suffix = ".unzip"
+
+    def _extract_file(self, fname, extract_dir):
+        """
+        This method receives an argument for the archive to extract and the
+        destination path.
+        """
+        with ZipFile(fname, "r") as zip_file:
+            if self.members is None:
+                warn("Unzipping contents of '{}' to '{}'".format(fname, extract_dir))
+                # Unpack all files from the archive into our new folder
+                zip_file.extractall(path=extract_dir)
+            else:
+                for member in self.members:
+                    warn(
+                        "Extracting '{}' from '{}' to '{}'".format(
+                            member, fname, extract_dir
+                        )
+                    )
+                    # Extract the data file from within the archive
+                    with zip_file.open(member) as data_file:
+                        # Save it to our desired file name
+                        with open(os.path.join(extract_dir, member), "wb") as output:
+                            output.write(data_file.read())
+
+
+class Untar(ExtractorProcessor):  # pylint: disable=too-few-public-methods
+    """
+    Processor that unpacks a tar archive and returns a list of all files.
+
+    Use with :meth:`pooch.Pooch.fetch` to untar a downloaded data file into a folder in
+    the local data store. :meth:`~pooch.Pooch.fetch` will return a list with the names
+    of the extracted files instead of the archive.
+
+    The output folder is ``{fname}.untar``.
+
+
+    Parameters
+    ----------
+    members : list or None
+        If None, will unpack all files in the archive. Otherwise, *members* must be
+        a list of file names to unpack from the archive. Only these files will be
+        unpacked.
+    """
+
+    suffix = ".untar"
+
+    def _extract_file(self, fname, extract_dir):
+        """
+        This method receives an argument for the archive to extract and the
+        destination path.
+        """
+        with TarFile.open(fname, "r") as tar_file:
+            if self.members is None:
+                warn("Untarring contents of '{}' to '{}'".format(fname, extract_dir))
+                # Unpack all files from the archive into our new folder
+                tar_file.extractall(path=extract_dir)
+            else:
+                for member in self.members:
+                    warn(
+                        "Extracting '{}' from '{}' to '{}'".format(
+                            member, fname, extract_dir
+                        )
+                    )
+                    # Extract the data file from within the archive
+                    # Python 2.7: extractfile doesn't return a context manager
+                    data_file = tar_file.extractfile(member)
+                    try:
+                        # Save it to our desired file name
+                        with open(os.path.join(extract_dir, member), "wb") as output:
+                            output.write(data_file.read())
+                    finally:
+                        data_file.close()


Test patch:
diff --git a/pooch/tests/data/registry-custom-url.txt b/pooch/tests/data/registry-custom-url.txt
index d68036e..97be99b 100644
--- a/pooch/tests/data/registry-custom-url.txt
+++ b/pooch/tests/data/registry-custom-url.txt
@@ -2,3 +2,5 @@ subdir/tiny-data.txt baee0894dba14b12085eacb204284b97e362f4f3e5a5807693cc90ef415
 tiny-data.zip 0d49e94f07bc1866ec57e7fd1b93a351fba36842ec9b13dd50bf94e8dfa35cbb
 store.zip 0498d2a001e71051bbd2acd2346f38da7cbd345a633cb7bf0f8a20938714b51a
 tiny-data.txt baee0894dba14b12085eacb204284b97e362f4f3e5a5807693cc90ef415c1b2d https://some-site/tiny-data.txt
+tiny-data.tar.gz 41503f083814f43a01a8e9a30c28d7a9fe96839a99727a7fdd0acf7cd5bab63b
+store.tar.gz 088c7f4e0f1859b1c769bb6065de24376f366374817ede8691a6ac2e49f29511
diff --git a/pooch/tests/data/registry.txt b/pooch/tests/data/registry.txt
index b339be9..8a7dcef 100644
--- a/pooch/tests/data/registry.txt
+++ b/pooch/tests/data/registry.txt
@@ -2,3 +2,5 @@ subdir/tiny-data.txt baee0894dba14b12085eacb204284b97e362f4f3e5a5807693cc90ef415
 tiny-data.txt baee0894dba14b12085eacb204284b97e362f4f3e5a5807693cc90ef415c1b2d
 tiny-data.zip 0d49e94f07bc1866ec57e7fd1b93a351fba36842ec9b13dd50bf94e8dfa35cbb
 store.zip 0498d2a001e71051bbd2acd2346f38da7cbd345a633cb7bf0f8a20938714b51a
+tiny-data.tar.gz 41503f083814f43a01a8e9a30c28d7a9fe96839a99727a7fdd0acf7cd5bab63b
+store.tar.gz 088c7f4e0f1859b1c769bb6065de24376f366374817ede8691a6ac2e49f29511
diff --git a/pooch/tests/test_core.py b/pooch/tests/test_core.py
index f1c1802..df44934 100644
--- a/pooch/tests/test_core.py
+++ b/pooch/tests/test_core.py
@@ -17,7 +17,6 @@ import pytest
 from .. import Pooch, create
 from ..utils import file_hash
 from ..downloaders import HTTPDownloader
-from ..processors import Unzip
 
 from .utils import pooch_test_url, pooch_test_registry, check_tiny_data
 
@@ -252,66 +251,6 @@ def test_check_availability():
     assert not pup.is_available("not-a-real-data-file.txt")
 
 
-def test_processors():
-    "Setup a post-download hook and make sure it's only executed when downloading"
-    unzip = Unzip(members=["tiny-data.txt"])
-    with TemporaryDirectory() as local_store:
-        path = Path(local_store)
-        true_path = str(path / "tiny-data.zip.unzip" / "tiny-data.txt")
-        # Setup a pooch in a temp dir
-        pup = Pooch(path=path, base_url=BASEURL, registry=REGISTRY)
-        # Check the warnings when downloading and from the processor
-        with warnings.catch_warnings(record=True) as warn:
-            fnames = pup.fetch("tiny-data.zip", processor=unzip)
-            fname = fnames[0]
-            assert len(fnames) == 1
-            assert len(warn) == 2
-            assert all(issubclass(w.category, UserWarning) for w in warn)
-            assert str(warn[-2].message).split()[0] == "Downloading"
-            assert str(warn[-1].message).startswith("Unzipping 'tiny-data.txt'")
-        assert fname == true_path
-        check_tiny_data(fname)
-        # Check that processor doesn't execute when not downloading
-        with warnings.catch_warnings(record=True) as warn:
-            fnames = pup.fetch("tiny-data.zip", processor=unzip)
-            fname = fnames[0]
-            assert len(fnames) == 1
-            assert not warn
-        assert fname == true_path
-        check_tiny_data(fname)
-
-
-def test_processor_multiplefiles():
-    "Setup a processor to unzip a file and return multiple fnames"
-    with TemporaryDirectory() as local_store:
-        path = Path(local_store)
-        true_paths = {
-            str(path / "store.zip.unzip" / "store" / "tiny-data.txt"),
-            str(path / "store.zip.unzip" / "store" / "subdir" / "tiny-data.txt"),
-        }
-        # Setup a pooch in a temp dir
-        pup = Pooch(path=path, base_url=BASEURL, registry=REGISTRY)
-        # Check the warnings when downloading and from the processor
-        with warnings.catch_warnings(record=True) as warn:
-            fnames = pup.fetch("store.zip", processor=Unzip())
-            assert len(warn) == 2
-            assert all(issubclass(w.category, UserWarning) for w in warn)
-            assert str(warn[-2].message).split()[0] == "Downloading"
-            assert str(warn[-1].message).startswith("Unzipping contents")
-            assert len(fnames) == 2
-            assert true_paths == set(fnames)
-            for fname in fnames:
-                check_tiny_data(fname)
-        # Check that processor doesn't execute when not downloading
-        with warnings.catch_warnings(record=True) as warn:
-            fnames = pup.fetch("store.zip", processor=Unzip())
-            assert not warn
-            assert len(fnames) == 2
-            assert true_paths == set(fnames)
-            for fname in fnames:
-                check_tiny_data(fname)
-
-
 def test_downloader():
     "Setup a downloader function for fetch"
 
diff --git a/pooch/tests/test_processors.py b/pooch/tests/test_processors.py
new file mode 100644
index 0000000..c2dae3f
--- /dev/null
+++ b/pooch/tests/test_processors.py
@@ -0,0 +1,109 @@
+"""
+Test the processor hooks
+"""
+from pathlib import Path
+
+try:
+    from tempfile import TemporaryDirectory
+except ImportError:
+    from backports.tempfile import TemporaryDirectory
+import warnings
+
+import pytest
+
+from .. import Pooch
+from ..processors import Unzip, Untar, ExtractorProcessor
+
+from .utils import pooch_test_url, pooch_test_registry, check_tiny_data
+
+
+REGISTRY = pooch_test_registry()
+BASEURL = pooch_test_url()
+
+
+def test_extractprocessor_fails():
+    "The base class should be used and should fail when passed to fecth"
+    with TemporaryDirectory() as local_store:
+        # Setup a pooch in a temp dir
+        pup = Pooch(path=Path(local_store), base_url=BASEURL, registry=REGISTRY)
+        processor = ExtractorProcessor()
+        with pytest.raises(NotImplementedError) as exception:
+            pup.fetch("tiny-data.tar.gz", processor=processor)
+        assert "'suffix'" in exception.value.args[0]
+        processor.suffix = "tar.gz"
+        with pytest.raises(NotImplementedError) as exception:
+            pup.fetch("tiny-data.tar.gz", processor=processor)
+        assert not exception.value.args
+
+
+@pytest.mark.parametrize(
+    "proc_cls,ext", [(Unzip, ".zip"), (Untar, ".tar.gz")], ids=["Unzip", "Untar"]
+)
+def test_processors(proc_cls, ext):
+    "Setup a post-download hook and make sure it's only executed when downloading"
+    processor = proc_cls(members=["tiny-data.txt"])
+    suffix = proc_cls.suffix
+    extract_dir = "tiny-data" + ext + suffix
+    with TemporaryDirectory() as local_store:
+        path = Path(local_store)
+        true_path = str(path / extract_dir / "tiny-data.txt")
+        # Setup a pooch in a temp dir
+        pup = Pooch(path=path, base_url=BASEURL, registry=REGISTRY)
+        # Check the warnings when downloading and from the processor
+        with warnings.catch_warnings(record=True) as warn:
+            fnames = pup.fetch("tiny-data" + ext, processor=processor)
+            fname = fnames[0]
+            assert len(fnames) == 1
+            assert len(warn) == 2
+            assert all(issubclass(w.category, UserWarning) for w in warn)
+            assert str(warn[-2].message).split()[0] == "Downloading"
+            assert str(warn[-1].message).startswith("Extracting 'tiny-data.txt'")
+        assert fname == true_path
+        check_tiny_data(fname)
+        # Check that processor doesn't execute when not downloading
+        with warnings.catch_warnings(record=True) as warn:
+            fnames = pup.fetch("tiny-data" + ext, processor=processor)
+            fname = fnames[0]
+            assert len(fnames) == 1
+            assert not warn
+        assert fname == true_path
+        check_tiny_data(fname)
+
+
+@pytest.mark.parametrize(
+    "proc_cls,ext,msg",
+    [(Unzip, ".zip", "Unzipping"), (Untar, ".tar.gz", "Untarring")],
+    ids=["Unzip", "Untar"],
+)
+def test_processor_multiplefiles(proc_cls, ext, msg):
+    "Setup a processor to unzip/untar a file and return multiple fnames"
+    processor = proc_cls()
+    suffix = proc_cls.suffix
+    extract_dir = "store" + ext + suffix
+    with TemporaryDirectory() as local_store:
+        path = Path(local_store)
+        true_paths = {
+            str(path / extract_dir / "store" / "tiny-data.txt"),
+            str(path / extract_dir / "store" / "subdir" / "tiny-data.txt"),
+        }
+        # Setup a pooch in a temp dir
+        pup = Pooch(path=path, base_url=BASEURL, registry=REGISTRY)
+        # Check the warnings when downloading and from the processor
+        with warnings.catch_warnings(record=True) as warn:
+            fnames = pup.fetch("store" + ext, processor=processor)
+            assert len(warn) == 2
+            assert all(issubclass(w.category, UserWarning) for w in warn)
+            assert str(warn[-2].message).split()[0] == "Downloading"
+            assert str(warn[-1].message).startswith("{} contents".format(msg))
+            assert len(fnames) == 2
+            assert true_paths == set(fnames)
+            for fname in fnames:
+                check_tiny_data(fname)
+        # Check that processor doesn't execute when not downloading
+        with warnings.catch_warnings(record=True) as warn:
+            fnames = pup.fetch("store" + ext, processor=processor)
+            assert not warn
+            assert len(fnames) == 2
+            assert true_paths == set(fnames)
+            for fname in fnames:
+                check_tiny_data(fname)
diff --git a/pooch/tests/utils.py b/pooch/tests/utils.py
index e25e948..9ea7f33 100644
--- a/pooch/tests/utils.py
+++ b/pooch/tests/utils.py
@@ -55,5 +55,7 @@ def pooch_test_registry():
         "subdir/tiny-data.txt": "baee0894dba14b12085eacb204284b97e362f4f3e5a5807693cc90ef415c1b2d",
         "tiny-data.zip": "0d49e94f07bc1866ec57e7fd1b93a351fba36842ec9b13dd50bf94e8dfa35cbb",
         "store.zip": "0498d2a001e71051bbd2acd2346f38da7cbd345a633cb7bf0f8a20938714b51a",
+        "tiny-data.tar.gz": "41503f083814f43a01a8e9a30c28d7a9fe96839a99727a7fdd0acf7cd5bab63b",
+        "store.tar.gz": "088c7f4e0f1859b1c769bb6065de24376f366374817ede8691a6ac2e49f29511",
     }
     return registry