CU-8699jpw1u Save and load CDB and Vocab as zip (#14)

* CU-8699jpw1u: Add utility method to save and load serialisables as .zip files.

And to do so automatically.

* CU-8699jpw1u: Allow a CDB to saved as a zip automatically.

* CU-8699jpw1u: A few simple tests for CDB automatic zip-serialisation

* CU-8699jpw1u: Add test time zipped CDB

* CU-8699jpw1u: Add option for vocab to be (automatically) saved as zip

* CU-8699jpw1u: A few simple tests for vocab zip-serialisation

* CU-8699jpw1u: Add test time zipped Vocab

Author: mart-r

medcat-v2/medcat/cdb/cdb.py

@@ -1,11 +1,13 @@
-from typing import Iterable, Any, Collection, Union
+from typing import Iterable, Any, Collection, Union, Literal
 
 from medcat.storage.serialisables import AbstractSerialisable
 from medcat.cdb.concepts import CUIInfo, NameInfo, TypeInfo
 from medcat.cdb.concepts import get_new_cui_info, get_new_name_info
 from medcat.cdb.concepts import reset_cui_training
 from medcat.storage.serialisers import (
     deserialise, AvailableSerialisers, serialise)
+from medcat.storage.zip_utils import (
+    should_serialise_as_zip, serialise_as_zip, deserialise_from_zip)
 from medcat.utils.defaults import default_weighted_average, StatusTypes as ST
 from medcat.utils.hasher import Hasher
 from medcat.preprocessors.cleaners import NameDescriptor
@@ -485,6 +487,7 @@ def save(self, save_path: str,
              serialiser: Union[
                  str, AvailableSerialisers] = AvailableSerialisers.dill,
              overwrite: bool = False,
+             as_zip: Union[bool, Literal['auto']] = 'auto',
              ) -> None:
         """Save CDB at path.
 
@@ -495,12 +498,20 @@ def save(self, save_path: str,
                 The serialiser. Defaults to AvailableSerialisers.dill.
             overwrite (bool, optional):
                 Whether to allow overwriting existing files. Defaults to False.
+            as_zip (Union[bool, Literal['auto']]):
+                Whether to serialise the CDB as a zip.
         """
+        if should_serialise_as_zip(save_path, as_zip):
+            serialise_as_zip(self, save_path, serialiser, overwrite=overwrite)
+            return
         serialise(serialiser, self, save_path, overwrite=overwrite)
 
     @classmethod
     def load(cls, path: str) -> 'CDB':
-        cdb = deserialise(path)
+        if should_serialise_as_zip(path, 'auto'):
+            cdb = deserialise_from_zip(path)
+        else:
+            cdb = deserialise(path)
         if not isinstance(cdb, CDB):
             raise ValueError(f"The path '{path}' is not a CDB!")
         return cdb

medcat-v2/medcat/storage/zip_utils.py

@@ -0,0 +1,69 @@
+import os
+import shutil
+import tempfile
+from typing import Union, Literal
+
+from medcat.storage.serialisables import Serialisable
+from medcat.storage.serialisers import (
+    serialise, deserialise, AvailableSerialisers)
+
+
+def should_serialise_as_zip(path: str,
+                            as_zip: Union[bool, Literal['auto']]
+                            ) -> bool:
+    if as_zip == 'auto':
+        return path.endswith(".zip")
+    return as_zip
+
+
+def serialise_as_zip(
+        obj: Serialisable, path: str,
+        ser_type: Union[AvailableSerialisers, str] = AvailableSerialisers.dill,
+        overwrite: bool = False,
+        ) -> None:
+    """Serialse the file to a .zip at the specified path.
+
+    The process uses the regular `serialise` method to serialise the object
+    as a folder into a temporary directory and subsequently zips it up to
+    the path requested.
+
+    Args:
+        obj (Serialisable): The object to serialise.
+        path (str): The path to serialse the file to. Should end with .zip.
+        ser_type (Union[AvailableSerialisers, str], optional): The serialiser
+            to use. Defaults to AvailableSerialisers.dill.
+        overwrite (bool, optional):
+            Whether to allow overwriting existing files. Defaults to False.
+    """
+    if not path.endswith('.zip'):
+        raise ValueError(f"Path must end with .zip, got {path}")
+    with tempfile.TemporaryDirectory() as tmpdir:
+        serialise(ser_type, obj, tmpdir)
+        if not overwrite and os.path.exists(path):
+            raise ValueError(f"Cannot overwrite existing file: {path}")
+        base_name = path[:-4]  # remove '.zip'
+        shutil.make_archive(
+            base_name=base_name,
+            format='zip',
+            root_dir=tmpdir,
+        )
+
+
+def deserialise_from_zip(path: str) -> Serialisable:
+    """Deserialise from a zip file.
+
+    The process involves unzipping the contents to a temporary directory,
+    and subsequently using the ruglar `deserialise` method to deserialise
+    from that.
+
+    Args:
+        path (str): The path to deserialise from. Should end with .zip.
+
+    Returns:
+        Serialisable: The deserialised object
+    """
+    if not path.endswith('.zip'):
+        raise ValueError(f"Path must end with .zip, got {path}")
+    with tempfile.TemporaryDirectory() as tmpdir:
+        shutil.unpack_archive(path, tmpdir, format='zip')
+        return deserialise(tmpdir)

medcat-v2/medcat/vocab.py

@@ -1,4 +1,4 @@
-from typing import Optional, Any, cast, Union
+from typing import Optional, Any, cast, Union, Literal
 from typing_extensions import TypedDict
 
 # import dill
@@ -7,6 +7,8 @@
 from medcat.storage.serialisables import AbstractSerialisable
 from medcat.storage.serialisers import (
     deserialise, AvailableSerialisers, serialise)
+from medcat.storage.zip_utils import (
+    should_serialise_as_zip, serialise_as_zip, deserialise_from_zip)
 
 
 WordDescriptor = TypedDict('WordDescriptor',
@@ -298,6 +300,7 @@ def save(self, save_path: str,
              serialiser: Union[
                  str, AvailableSerialisers] = AvailableSerialisers.dill,
              overwrite: bool = False,
+             as_zip: Union[bool, Literal['auto']] = 'auto',
              ) -> None:
         """Save Vocab at path.
 
@@ -308,12 +311,20 @@ def save(self, save_path: str,
                 The serialiser. Defaults to AvailableSerialisers.dill.
             overwrite (bool, optional):
                 Whether to allow overwriting existing files. Defaults to False.
+            as_zip (Union[bool, Literal['auto']]):
+                Whether to serialise the CDB as a zip.
         """
+        if should_serialise_as_zip(save_path, as_zip):
+            serialise_as_zip(self, save_path, serialiser, overwrite=overwrite)
+            return
         serialise(serialiser, self, save_path, overwrite=overwrite)
 
     @classmethod
     def load(cls, path: str) -> 'Vocab':
-        vocab = deserialise(path)
+        if should_serialise_as_zip(path, 'auto'):
+            vocab = deserialise_from_zip(path)
+        else:
+            vocab = deserialise(path)
         if not isinstance(vocab, Vocab):
             raise ValueError(f"The path '{path}' is not a Vocab!")
         return vocab

medcat-v2/tests/cdb/test_cdb.py

@@ -9,7 +9,10 @@
 from unittest import TestCase
 import tempfile
 
-from .. import UNPACKED_EXAMPLE_MODEL_PACK_PATH
+from .. import UNPACKED_EXAMPLE_MODEL_PACK_PATH, RESOURCES_PATH
+
+
+ZIPPED_CDB_PATH = os.path.join(RESOURCES_PATH, "mct2_cdb.zip")
 
 
 class CDBTests(TestCase):
@@ -26,9 +29,28 @@ def test_convenience_method_save(self):
         with tempfile.TemporaryDirectory() as dir:
             self.cdb.save(dir)
             self.assertTrue(os.path.exists(dir))
+            # should have a non-empty directory
+            self.assertTrue(os.listdir(dir))
             obj = deserialise(dir)
             self.assertIsInstance(obj, cdb.CDB)
 
+    def test_can_load_from_zip(self):
+        loaded = cdb.CDB.load(ZIPPED_CDB_PATH)
+        self.assertIsInstance(loaded, cdb.CDB)
+        # make sure it's actually a file not a folder
+        self.assertTrue(os.path.isfile(ZIPPED_CDB_PATH))
+
+    def test_can_save_to_zip(self):
+        with tempfile.TemporaryDirectory() as temp_dir:
+            file_name = os.path.join(temp_dir, "cdb.zip")
+            # NOTE: auto detection should write as zip
+            self.cdb.save(file_name)
+            self.assertTrue(os.path.exists(file_name))
+            self.assertTrue(os.path.isfile(file_name))
+            # and can load from saved zip
+            loaded = cdb.CDB.load(file_name)
+            self.assertIsInstance(loaded, cdb.CDB)
+
     def test_convenience_method_load(self):
         ccdb = cdb.CDB.load(self.CDB_PATH)
         self.assertIsInstance(ccdb, cdb.CDB)

medcat-v2/tests/resources/mct2_cdb.zip

@@ -8,7 +8,10 @@
 import unittest
 import tempfile
 
-from . import UNPACKED_EXAMPLE_MODEL_PACK_PATH
+from . import UNPACKED_EXAMPLE_MODEL_PACK_PATH, RESOURCES_PATH
+
+
+ZIPPED_VOCAB_PATH = os.path.join(RESOURCES_PATH, "mct2_vocab.zip")
 
 
 class VocabCreationTests(unittest.TestCase):
@@ -191,3 +194,18 @@ def test_convenience_save(self):
     def test_convenience_load(self):
         vocab = Vocab.load(self.VOCAB_PATH)
         self.assertIsInstance(vocab, Vocab)
+
+    def test_can_load_from_zip(self):
+        vocab = Vocab.load(ZIPPED_VOCAB_PATH)
+        self.assertIsInstance(vocab, Vocab)
+
+    def test_can_save_to_zip(self):
+        with tempfile.TemporaryDirectory() as temp_dir:
+            file_name = os.path.join(temp_dir, 'vocab.zip')
+            # NOTE: auto detection will write as zip
+            self.vocab.save(file_name)
+            self.assertTrue(os.path.exists(file_name))
+            self.assertTrue(os.path.isfile(file_name))
+            # and can load from saved zip
+            loaded = Vocab.load(file_name)
+            self.assertIsInstance(loaded, Vocab)