/ src / python / txtai / archive / base.py
base.py
  1  """
  2  Archive module
  3  """
  4  
  5  import os
  6  
  7  from tempfile import TemporaryDirectory
  8  
  9  from .tar import Tar
 10  from .zip import Zip
 11  
 12  
 13  class Archive:
 14      """
 15      Base class for archive instances.
 16      """
 17  
 18      def __init__(self, directory=None):
 19          """
 20          Creates a new archive instance.
 21  
 22          Args:
 23              directory: directory to use as working directory, defaults to a temporary directory
 24          """
 25  
 26          self.directory = directory
 27  
 28      def isarchive(self, path):
 29          """
 30          Checks if path is an archive file based on the extension.
 31  
 32          Args:
 33              path: path to check
 34  
 35          Returns:
 36              True if the path ends with an archive extension, False otherwise
 37          """
 38  
 39          return path and any(path.lower().endswith(extension) for extension in [".tar.bz2", ".tar.gz", ".tar.xz", ".zip"])
 40  
 41      def path(self):
 42          """
 43          Gets the current working directory for this archive instance.
 44  
 45          Returns:
 46              archive working directory
 47          """
 48  
 49          # Default to a temporary directory. All files created in this directory will be deleted
 50          # when this archive instance goes out of scope.
 51          if not self.directory:
 52              # pylint: disable=R1732
 53              self.directory = TemporaryDirectory()
 54  
 55          return self.directory.name if isinstance(self.directory, TemporaryDirectory) else self.directory
 56  
 57      def load(self, path, compression=None):
 58          """
 59          Extracts file at path to archive working directory.
 60  
 61          Args:
 62              path: path to archive file
 63              compression: compression format, infers from path if not provided
 64          """
 65  
 66          # Unpack compressed file
 67          compress = self.create(path, compression)
 68          compress.unpack(path, self.path())
 69  
 70      def save(self, path, compression=None):
 71          """
 72          Archives files in archive working directory to file at path.
 73  
 74          Args:
 75              path: path to archive file
 76              compression: compression format, infers from path if not provided
 77          """
 78  
 79          # Create output directory, if necessary
 80          output = os.path.dirname(path)
 81          if output:
 82              os.makedirs(output, exist_ok=True)
 83  
 84          # Pack into compressed file
 85          compress = self.create(path, compression)
 86          compress.pack(self.path(), path)
 87  
 88      def create(self, path, compression):
 89          """
 90          Method to construct a Compress instance.
 91  
 92          Args:
 93              path: file path
 94              compression: compression format, infers using file extension if not provided
 95  
 96          Returns:
 97              Compress
 98          """
 99  
100          # Infer compression format from path if not provided
101          compression = compression if compression else path.lower().split(".")[-1]
102  
103          # Create compression instance
104          return Zip() if compression == "zip" else Tar()