Unify FileCorpus classes to incorporate archives (zip/tar/gz/...)

Author: robert-lieck

corpusinterface/corpora.py

@@ -1,53 +1,78 @@
-#  Copyright (c) 2020 Robert Lieck
 from pathlib import Path
+from abc import ABC, abstractmethod
 import os
+import io
 import re
 import json
 import ast
+import zipfile
+import tarfile
 
 import pandas
 
 
-class Data:
+class Corpus(ABC):
     """An abstract base class for data items, such as corpora or documents."""
+
+    def __enter__(self):
+        self.open()
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.close()
+        return False
+
+    def open(self):
+        """Open resources owned by the corpus."""
+        pass
+
+    def close(self):
+        """Close resources owned by the corpus."""
+        pass
+
     def metadata(self, *args, **kwargs):
-        raise NotImplementedError
+        return None
 
+    @abstractmethod
     def data(self, *args, **kwargs):
         raise NotImplementedError
 
 
-class FileCorpus(Data):
-    """A collection of files in a directory"""
+class FileCorpusBase(Corpus):
+    """A collection of file-like objects in a directory-like object (e.g., normal files or a zip archive)."""
 
     # special keyword arguments
     __META_READER__ = "meta_reader"
     __FILE_READER__ = "file_reader"
 
-    @classmethod
-    def init(cls, **kwargs):
-        if 'path' not in kwargs:
-            raise TypeError("Missing required keyword argument 'path'")
-        kwargs = {**dict(file_regex=None,
-                         path_regex=None,
-                         file_exclude_regex=None,
-                         path_exclude_regex=None),
-                  **kwargs}
-        return FileCorpus(**kwargs)
-
-    def __init__(self,
+    class LazyLoad:
+        """Simple wrapper for lazy-loading of data using a provided loading function."""
+
+        def __init__(self, load_func, *args, **kwargs):
+            self.load_func = load_func
+            self.args = args
+            self.kwargs = kwargs
+            self.data = None
+
+        def load(self):
+            if self.data is None:
+                self.data = self.load_func(*self.args, **self.kwargs)
+            return self.data
+
+    def __init__(self, *,
                  path,
+                 include_dirs=False,
                  file_regex=None,
                  path_regex=None,
                  file_exclude_regex=None,
                  path_exclude_regex=None,
                  **kwargs):
         # set path and check
-        self.path = Path(path)
-        if not self.path.exists():
-            raise FileNotFoundError(f"Corpus directory {self.path} does not exist")
-        elif not self.path.is_dir():
-            raise NotADirectoryError(f"{self.path} is not a directory")
+        self._path = Path(path)
+        if not self._path.exists():
+            raise FileNotFoundError(f"Corpus path '{self._path}' does not exist")
+        # whether to include directories when iterating over the corpus
+        self.include_dirs = include_dirs
         # remember additional keyword arguments
         self.kwargs = kwargs
         # initialise regex for including files
@@ -72,50 +97,265 @@ def __init__(self,
             self.path_exclude_regex = re.compile(path_exclude_regex)
 
     def __repr__(self):
-        return f"{self.__class__.__name__}({self.path})"
+        return f"{self.__class__.__name__}({self._path})"
+
+    def _skip_file(self, file_name, file_path):
+        # check file inclusion regex
+        if self.file_regex is not None and not self.file_regex.match(file_name):
+            return True  # skip non-matching files
+        # check path inclusion regex
+        if self.path_regex is not None and not self.path_regex.match(file_path):
+            return True  # skip non-matching paths
+        # check file exclusion regex
+        if self.file_exclude_regex is not None and self.file_exclude_regex.match(file_name):
+            return True  # skip matching files
+        # check path exclusion regex
+        if self.path_exclude_regex is not None and self.path_exclude_regex.match(file_path):
+            return True  # skip matching paths
+        return False
+
+    @abstractmethod
+    def files(self):
+        """Return an iterator over the files in the corpus."""
+        raise NotImplementedError
 
     def metadata(self, *args, **kwargs):
         kwargs = {**self.kwargs, **kwargs}
         meta_reader = kwargs.get(self.__META_READER__, None)
         if meta_reader is None:
-            return self.path
+            return self._path
         else:
-            return meta_reader(self.path, *args, **kwargs)
+            return meta_reader(self._path, *args, **kwargs)
+
+    def data(self, *args, **kwargs):
+        kwargs = {**dict(return_files=False, lazy_load=False), **self.kwargs, **kwargs}
+        return_files = kwargs.pop('return_files')
+        lazy_load = kwargs.pop('lazy_load')
+        file_reader = kwargs.get(self.__FILE_READER__, None)
+        if file_reader is None:
+            raise TypeError(f"{self.__FILE_READER__} is None; provide file_reader or use files() if you want to iterate over the files instead")
+        if not callable(file_reader):
+            raise TypeError(f"{self.__FILE_READER__} must be a callable, not '{file_reader}' of type {type(file_reader)}")
+        for file_path in self.files():
+            if lazy_load:
+                d = self.LazyLoad(file_reader, file_path, *args, **kwargs)
+            else:
+                d = file_reader(file_path, *args, **kwargs)
+            if return_files:
+                yield file_path, d
+            else:
+                yield d
+
+
+class PathLike(ABC):
+    """
+    An abstract base class for path-like objects that behave roughly like pathlib.Path. They are NOT guaranteed to
+    exist as real paths in the filesystem (e.g., they could be files within a zip archive) and therefore do not
+    implement the os.PathLike API.
+    """
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}('{self.path}')"
+
+    def __str__(self):
+        return str(self.path)
+
+    @property
+    @abstractmethod
+    def name(self):
+        """Name of the file or directory the path points to."""
+        raise NotImplementedError
+
+    @property
+    @abstractmethod
+    def path(self):
+        """Full path to the file or directory the path points to."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def open(self, mode='r'):
+        """Open the file the path points to."""
+        raise NotImplementedError
+
+
+class FileCorpus(FileCorpusBase):
+    """A collection of normal files in a directory"""
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        # check if path is a directory
+        if not self._path.is_dir():
+            raise NotADirectoryError(f"{self._path} is not a directory")
 
     def files(self):
         # recursively traverse directory
-        for root, dirs, files in os.walk(self.path):
+        for root, dirs, files in os.walk(self._path):
             root = Path(root)
-            for file in files:
-                path = root / file
-                # check file inclusion regex
-                if self.file_regex is not None and not self.file_regex.match(file):
-                    continue # skip non-matching files
-                # check path inclusion regex
-                if self.path_regex is not None and not self.path_regex.match(str(path)):
-                    continue  # skip non-matching paths
-                # check file exclusion regex
-                if self.file_exclude_regex is not None and self.file_exclude_regex.match(file):
-                    continue  # skip matching files
-                # check path exclusion regex
-                if self.path_exclude_regex is not None and self.path_exclude_regex.match(str(path)):
-                    continue  # skip matching paths
-                # yield absolute path to file
-                yield path
+            for file_name in files:
+                file_path = root / file_name
+                if self._skip_file(file_name, str(file_path)):
+                    continue
+                if file_path.is_dir() and not self.include_dirs:
+                    continue
+                yield RealPath(file_path)
+
+
+class RealPath(PathLike, os.PathLike):
+    """
+    Class representing a real path. This is a subclass of PathLike and os.PathLike to make it compatible with both
+    standard os.PathLike objects and virtual paths in archive files (e.g., ZipFile, TarFile).
+    """
+
+    def __init__(self, path):
+        self._path = Path(path)
+
+    def __fspath__(self):
+        return os.fspath(self._path)
+
+    @property
+    def name(self):
+        return self._path.name
+
+    @property
+    def path(self):
+        return self._path
+
+    def open(self, mode='r'):
+        return self._path.open(mode)
+
+
+class ArchiveCorpusBase(FileCorpusBase):
+    def __init__(self, path, file=None, **kwargs):
+        path = Path(path)
+        if file is not None:
+            path /= Path(file)
+        super().__init__(path=path, **kwargs)
+        if not self._path.is_file():
+            raise FileNotFoundError(f"'{self._path}' is not an archive file")
+        self._archive = None
+
+
+class ZipFileCorpus(ArchiveCorpusBase):
+
+    def open(self) -> None:
+        if self._archive is not None:
+            raise RuntimeError("Archive is already open")
+        self._archive = zipfile.ZipFile(self._path, mode="r")
+
+    def close(self) -> None:
+        if self._archive is not None:
+            self._archive.close()
+            self._archive = None
+
+    def files(self):
+        if self._archive is None:
+            raise RuntimeError("Archive must be opened before iteration")
+        for info in self._archive.infolist():
+            file_name = Path(info.filename).name
+            file_path = self._path / Path(info.filename)
+
+            if self._skip_file(str(file_name), str(file_path)):
+                continue
+
+            if info.is_dir() and not self.include_dirs:
+                continue
+
+            yield ZipEntry(self, info)
+
+
+class ZipEntry(PathLike):
+
+    def __init__(self, archive, info):
+        self._archive = archive
+        self._info = info
+
+    @property
+    def name(self):
+        return Path(self._info.filename).name
+
+    @property
+    def path(self):
+        return self._archive._path / Path(self._info.filename)
+
+    def open(self, mode="r", **kwargs):
+        binary = self._archive._archive.open(self._info, "r")
+        if mode == "rb":
+            return binary
+        elif mode in ("r", "rt"):
+            kwargs.setdefault('encoding', "utf-8")
+            return io.TextIOWrapper(binary, **kwargs)
+        else:
+            raise ValueError(f"Unsupported mode: {mode}")
+
+
+class TarFileCorpus(ArchiveCorpusBase):
+
+    def open(self):
+        if self._archive is not None:
+            raise RuntimeError("Archive is already open")
+
+        # Automatically detects .tar, .tar.gz, .tgz, .tar.bz2, .tar.xz, etc.
+        self._archive = tarfile.open(self._path, mode="r:*")
+
+    def close(self):
+        if self._archive is not None:
+            self._archive.close()
+            self._archive = None
+
+    def files(self):
+        if self._archive is None:
+            raise RuntimeError("Archive must be opened before iteration")
+
+        for info in self._archive.getmembers():
+            file_name = Path(info.name).name
+            file_path = self._path / Path(info.name)
+
+            if self._skip_file(str(file_name), str(file_path)):
+                continue
+
+            if info.isdir() and not self.include_dirs:
+                continue
+
+            # Skip entries such as symbolic links and device files.
+            if not info.isfile() and not info.isdir():
+                continue
+
+            yield TarEntry(self, info)
+
+
+class TarEntry(PathLike):
+
+    def __init__(self, archive, info):
+        self._archive = archive
+        self._info = info
+
+    @property
+    def name(self):
+        return Path(self._info.name).name
+
+    @property
+    def path(self):
+        return self._archive._path / Path(self._info.name)
+
+    def open(self, mode="r", **kwargs):
+        binary = self._archive._archive.extractfile(self._info)
+
+        if binary is None:
+            raise OSError(f"Could not open archive entry: {self.path}")
+
+        if mode == "rb":
+            return binary
+        elif mode in ("r", "rt"):
+            kwargs.setdefault('encoding', "utf-8")
+            return io.TextIOWrapper(binary, **kwargs)
+        else:
+            raise ValueError(f"Unsupported mode: {mode}")
 
-    def data(self, *args, **kwargs):
-        kwargs = {**self.kwargs, **kwargs}
-        file_reader = kwargs.get(self.__FILE_READER__, None)
-        for path in self.files():
-            if file_reader is None:
-                yield path
-            else:
-                yield file_reader(path, *args, **kwargs)
 
 # single file corpora
 # -------------------
 
-class SingleFileCorpus(Data):
+class SingleFileCorpus(Corpus):
     """
     A corpus consisting of a single file, superclass for different file types.