Greatly simplify the template saving logic

Author: blythed

superduper/__init__.py

@@ -35,7 +35,7 @@
 from .components.plugin import Plugin
 from .components.streamlit import Streamlit
 from .components.table import Table
-from .components.template import QueryTemplate, Template
+from .components.template import Template
 from .components.vector_index import VectorIndex
 
 REQUIRES = [
@@ -65,7 +65,6 @@
     'Table',
     'Application',
     'Template',
-    'QueryTemplate',
     'Application',
     'Component',
     'pickle_serializer',

superduper/base/base.py

@@ -201,7 +201,7 @@ def cls_encode(cls, item: 'Base', builds, blobs, files, leaves_to_keep=()):
         )
 
     @staticmethod
-    def decode(r, db: t.Optional['Datalayer'] = None):
+    def decode(r):
         """Decode a dictionary component into a `Component` instance.
 
         :param r: Object to be decoded.
@@ -211,9 +211,8 @@ def decode(r, db: t.Optional['Datalayer'] = None):
 
         if '_path' in r:
             from superduper.misc.importing import import_object
-
             cls = import_object(r['_path'])
-        r = Document.decode(r, schema=cls.class_schema, db=db)
+        r = Document.decode(r, schema=cls.class_schema)
         return cls.from_dict(r, db=None)
 
     def encode(

superduper/base/datatype.py

@@ -298,6 +298,35 @@ def hash(cls, item):
         return hash_item([x.hash for x in item])
 
 
+class FDict(BaseDataType):
+    dtype: t.ClassVar[str] = 'json'
+
+    def encode_data(self, item, context):
+        """Encode the given item into a bytes-like object or reference.
+
+        :param item: The object/instance to encode.
+        :param context: A context object containing caches.
+        """
+        assert isinstance(item, dict)
+        return {
+            k: File().encode_data(v, context)
+            for k, v in item.items()
+        }
+
+    def decode_data(self, item, builds, db):
+        """Decode the item from `bytes`.
+
+        :param item: The item to decode.
+        :param builds: The builds.
+        :param db: The Datalayer.
+        """
+        return {k: File().decode_data(v, builds, db) for k, v in item.items()}
+
+    @classmethod
+    def hash(cls, item):
+        return hash_item({k: File.hash(v) for k, v in item.items()})
+
+
 @dc.dataclass(kw_only=True)
 class BaseVector(BaseDataType):
     """Base class for vector.
@@ -453,8 +482,11 @@ def hash_indescript(item):
     """
     if inspect.isfunction(item):
         module = item.__module__
-        body = f'{module}\n{inspect.getsource(item)}'
-        return hashlib.sha256(body.encode()).hexdigest()
+        try:
+            body = f'{module}\n{inspect.getsource(item)}'
+            return hashlib.sha256(body.encode()).hexdigest()
+        except OSError:
+            return hashlib.sha256(str(item).encode()).hexdigest()
     if inspect.isclass(item):
         module = item.__module__
         body = f'{module}\n{inspect.getsource(item)}'
@@ -601,6 +633,8 @@ def encode_data(self, item, context):
         :param item: The object/instance to encode.
         :param context: A context object containing caches.
         """
+        if isinstance(item, FileItem):
+            return item.reference
         assert os.path.exists(item)
         file = FileItem(identifier=self.hash(item), path=item)
         context.files[file.identifier] = file.path
@@ -690,6 +724,7 @@ class _DatatypeLookup:
             File(),
             LeafType(),
             ComponentType(),
+            FDict(),
             SDict(),
             SList(),
             FieldType('str'),

superduper/base/schema.py

@@ -1,10 +1,11 @@
 # TODO move to base
 import dataclasses as dc
 import json
+import re
 import typing as t
 from functools import cached_property
 
-from superduper import CFG
+from superduper import CFG, logging
 from superduper.base.datatype import BaseDataType
 from superduper.base.encoding import EncodeContext
 from superduper.misc.special_dicts import dict_to_ascii_table
@@ -40,14 +41,44 @@ def __add__(self, other: 'Schema'):
         new_fields.update(other.fields)
         return Schema(fields=new_fields)
 
-    @cached_property
+    @property
     def trivial(self):
         """Determine if the schema contains only trivial fields."""
         return not any([isinstance(v, BaseDataType) for v in self.fields.values()])
 
     def __repr__(self):
         return dict_to_ascii_table(self.fields)
 
+    @staticmethod
+    def handle_references(item, builds):
+        if '?(' not in str(item):
+            return item
+
+        if isinstance(item, str):
+            instances = re.findall(r'\?\((.*?)\)', item)
+
+            for k in instances:
+                name = k.split('.')[0]
+                attr = k.split('.')[-1]
+
+                if name not in builds:
+                    logging.warn(f'Could not find reference {name} from reference in {item} in builds')
+                    return item
+
+                to_replace = getattr(builds[name], attr)
+                item = item.replace(f'?({k})', str(to_replace))
+
+            return item
+        elif isinstance(item, list):
+            return [Schema.handle_references(i, builds) for i in item]
+        elif isinstance(item, dict):
+            return {
+                Schema.handle_references(k, builds): Schema.handle_references(v, builds)
+                for k, v in item.items()
+            }
+        else:
+            return item
+
     def decode_data(
         self, data: dict[str, t.Any], builds: t.Dict, db
     ) -> dict[str, t.Any]:
@@ -62,9 +93,17 @@ def decode_data(
 
         decoded = {}
 
+        # reorder the component so that references go first
+        is_ref = lambda x: isinstance(x, str) and x.startswith('?')
+        data_is_ref = {k: v for k, v in data.items() if is_ref(v)}
+        data_not_ref = {k: v for k, v in data.items() if not is_ref(v)}
+        data = {**data_is_ref, **data_not_ref}
+
         for k, value in data.items():
             field = self.fields.get(k)
 
+            value = self.handle_references(value, builds)
+
             if not isinstance(field, BaseDataType) or value is None:
                 decoded[k] = value
                 continue
@@ -133,6 +172,9 @@ def encode_data(self, out, context: t.Optional[EncodeContext] = None, **kwargs):
 
         return result
 
+    def __getitem__(self, item: str):
+        return self.fields[item]
+
 
 def get_schema(db, schema: t.Union[Schema, str]) -> t.Optional[Schema]:
     """Handle schema caching and loading.

superduper/components/component.py

@@ -56,8 +56,8 @@ def _build_info_from_path(path: str):
         files = {}
         for file_id in os.listdir(os.path.join(path, "files")):
             sub_paths = os.listdir(os.path.join(path, "files", file_id))
-            assert len(sub_paths) == 1, f"Multiple files found in {file_id}"
-            file_name = sub_paths[0]
+            # assert len(sub_paths) == 1, f"Multiple files found in {file_id}"
+            file_name = next(x for x in sub_paths if not x.startswith(".") or x.startswith("_"))
             files[file_id] = os.path.join(path, "files", file_id, file_name)
         config_object[KEY_FILES] = files
 
@@ -484,19 +484,16 @@ def dependencies(self):
         """Get dependencies on the component."""
         return ()
 
-    # TODO why both methods?
     def init(self):
         """Method to help initiate component field dependencies."""
-        self.unpack()
 
-    def unpack(self):
-        """Method to unpack the component.
-
-        This method is used to initialize all the fields of the component and leaf
-        """
+        def mro(item):
+            objects = item.__class__.__mro__
+            return [f'{o.__module__}.{o.__name__}' for o in objects]
 
         def _init(item):
-            if isinstance(item, Component):
+            
+            if 'superduper.components.component.Component' in mro(item):
                 item.init()
                 return item
 
@@ -552,7 +549,7 @@ def declare_component(self, cluster):
         pass
 
     @staticmethod
-    def read(path: str, db: t.Optional['Datalayer'] = None):
+    def read(path: str):
         """
         Read a `Component` instance from a directory created with `.export`.
 
@@ -566,45 +563,8 @@ def read(path: str, db: t.Optional['Datalayer'] = None):
         |_files/*
         ```
         """
-        was_zipped = False
-        if path.endswith('.zip'):
-            was_zipped = True
-            import shutil
-
-            shutil.unpack_archive(path)
-            path = path.replace('.zip', '')
-
         config_object = _build_info_from_path(path=path)
-
-        from superduper import Document
-
-        if db is not None:
-            for blob in os.listdir(path + '/' + 'blobs'):
-                with open(path + '/blobs/' + blob, 'rb') as f:
-                    data = f.read()
-                    db.artifact_store.put_bytes(data, blob)
-
-            out = Document.decode(config_object, db=db).unpack()
-        else:
-            from superduper.base.artifacts import FileSystemArtifactStore
-
-            artifact_store = FileSystemArtifactStore(
-                conn=path,
-                name='tmp_artifact_store',
-                files='files',
-                blobs='blobs',
-            )
-            db = namedtuple('tmp_db', field_names=('artifact_store',))(
-                artifact_store=artifact_store
-            )
-            cls = import_object(config_object['_path'])
-            out = Document.decode(
-                config_object, schema=cls.class_schema, db=db
-            ).unpack()
-            out = cls.from_dict(out, db=db)
-            if was_zipped:
-                shutil.rmtree(path)
-        return out
+        return Component.decode(config_object)
 
     def export(
         self,

superduper/components/plugin.py

@@ -23,11 +23,11 @@ class Plugin(Component):
 
     def postinit(self):
         """Post initialization method."""
-        if isinstance(self.path, FileItem):
-            self._prepare_plugin()
-        else:
-            path_name = os.path.basename(self.path.rstrip("/"))
-            self.identifier = self.identifier or f"plugin-{path_name}".replace(".", "_")
+        self.init()
+        self._prepare_plugin()
+
+        path_name = os.path.basename(self.path.rstrip("/"))
+        self.identifier = self.identifier or f"plugin-{path_name}".replace(".", "_")
         self._install()
         super().postinit()
 
@@ -91,19 +91,22 @@ def _pip_install(self, requirement_path):
 
     def _prepare_plugin(self):
         plugin_name_tag = f"{self.identifier}"
-        assert isinstance(self.path, FileItem)
+        if isinstance(self.path, FileItem):
+            self.path = self.path.unpack()
+
         cache_path = os.path.expanduser(self.cache_path)
         uuid_path = os.path.join(cache_path, self.uuid)
+
         # Check if plugin is already in cache
         if os.path.exists(uuid_path):
             names = os.listdir(uuid_path)
             names = [name for name in names if name != "__pycache__"]
             assert len(names) == 1, f"Multiple plugins found in {uuid_path}"
             self.path = os.path.join(uuid_path, names[0])
+            sys.path.append(uuid_path)
             return
 
         logging.info(f"Preparing plugin {plugin_name_tag}")
-        self.path = self.path.unpack()
         assert os.path.exists(
             self.path
         ), f"Plugin {plugin_name_tag} not found at {self.path}"