chore(storage): abc improvements, initial Archive work

Author: gadicc

api/utils/storage/BaseStorage.py

@@ -0,0 +1,97 @@
+import os
+import re
+import subprocess
+from abc import ABC, abstractmethod
+
+
+class BaseArchive(ABC):
+    def __init__(self, path):
+        self.path = path
+
+    def extract(self):
+        print("TODO")
+
+    def splitext(self):
+        base, ext = os.path.splitext(self.path)
+        base, subext = os.path.splitext(base)
+        return base, ext, subext
+
+
+class TarZstdArchive(BaseArchive):
+    @staticmethod
+    def test(path):
+        return re.search(r"\.tar\.zstd?$", path)
+
+    def extract(self, dir, dry_run=False):
+        if not dir:
+            dir = os.path.dirname(self.path)
+        base, ext, subext = self.splitext()
+        dir = os.path.join(dir, base)
+
+        if not dry_run:
+            os.mkdir(dir)
+            subprocess.run(
+                [
+                    "tar",
+                    "--use-compress-program=unzstd",
+                    "-C",
+                    dir,
+                    "-xvf",
+                    self.path,
+                ],
+                check=True,
+            )
+            os.remove(self.path)
+
+        return dir  # , base, ext, subext
+
+
+archiveClasses = [TarZstdArchive]
+
+
+def Archive(path, **kwargs):
+    for ArchiveClass in archiveClasses:
+        if ArchiveClass.test(path):
+            return ArchiveClass(path, **kwargs)
+
+
+class BaseStorage(ABC):
+    @staticmethod
+    @abstractmethod
+    def test(url):
+        return re.search(r"^https?://", url)
+
+    def __init__(self, url, **kwargs):
+        self.url = url
+
+    def splitext(self):
+        base, ext = os.path.splitext(self.url)
+        base, subext = os.path.splitext(base)
+        return base, ext, subext
+
+    def get_filename(self):
+        return self.url.split("/").pop()
+
+    @abstractmethod
+    def download_file(self, dest):
+        """Download the file to `dest`"""
+        pass
+
+    def download_and_extract(self, fname, dry_run=False):
+        """
+        Downloads the file, and if it's an archive, extract it too.  Returns
+        the filename if not, or directory name (fname without extension) if
+        it was.
+        """
+        if not fname:
+            fname = self.get_filename()
+
+        dir = None
+        archive = Archive(fname, dry_run=dry_run)
+        if archive:
+            # TODO, streaming pipeline
+            self.download_file(fname)
+            return archive.extract()
+        else:
+            self.download_file(fname)
+            return fname

api/utils/storage/BaseStorage_test.py

@@ -0,0 +1,22 @@
+import unittest
+from . import Storage, S3Storage, HTTPStorage
+
+
+class BaseStorageTest(unittest.TestCase):
+    def test_get_filename(self):
+        storage = Storage("http://host.com/dir/file.tar.zst")
+        self.assertEqual(storage.get_filename(), "file.tar.zst")
+
+    class Download_and_extract(unittest.TestCase):
+        def test_file_only(self):
+            storage = Storage("http://host.com/dir/file.bin")
+            result = storage.download_and_extract(dry_run=True)
+            self.assertEqual(result, "file.bin")
+
+        def test_file_archive(self):
+            storage = Storage("http://host.com/dir/file.tar.zst")
+            result, base, ext, subext = storage.download_and_extract(dry_run=True)
+            self.assertEqual(result, "file")
+            self.assertEqual(base, "file")
+            self.assertEqual(ext, "tar")
+            self.assertEqual(subext, "zst")

api/utils/storage/HTTPStorage.py

@@ -3,13 +3,14 @@
 import time
 import requests
 from tqdm import tqdm
+from .BaseStorage import BaseStorage
 
 
 def get_now():
     return round(time.time() * 1000)
 
 
-class HTTPStorage:
+class HTTPStorage(BaseStorage):
     @staticmethod
     def test(url):
         return re.search(r"^https?://", url)

api/utils/storage/S3Storage.py

@@ -5,7 +5,7 @@
 import time
 from tqdm import tqdm
 from botocore.client import Config
-
+from .BaseStorage import BaseStorage
 
 AWS_S3_ENDPOINT_URL = os.environ.get("AWS_S3_ENDPOINT_URL", None)
 AWS_S3_DEFAULT_BUCKET = os.environ.get("AWS_S3_DEFAULT_BUCKET", None)
@@ -19,7 +19,7 @@ def get_now():
     return round(time.time() * 1000)
 
 
-class S3Storage:
+class S3Storage(BaseStorage):
     def test(url):
         return re.search(r"^(https?\+)?s3://", url)