Register artifacts in wrapper constructors

vitessce/wrappers.py

@@ -47,6 +47,7 @@ def __init__(self, **kwargs):
         self.file_def_creators = []
         self.base_dir = None
         self.stores = {}
+        self.artifacts = {}
         self._request_init = kwargs['request_init'] if 'request_init' in kwargs else None
 
     def __repr__(self):
@@ -73,6 +74,14 @@ def get_routes(self):
         :rtype: list[starlette.routing.Route]
         """
         return self.routes
+
+    def register_artifact(self, artifact):
+        artifact_url = artifact.path.to_url()
+        self.artifacts[artifact_url] = artifact
+        return artifact_url
+
+    def get_artifacts(self):
+        return self.artifacts
 
     def get_stores(self, base_url):
         """
@@ -417,21 +426,37 @@ class ImageOmeTiffWrapper(AbstractWrapper):
     :param \\*\\*kwargs: Keyword arguments inherited from :class:`~vitessce.wrappers.AbstractWrapper`
     """
 
-    def __init__(self, img_path=None, offsets_path=None, img_url=None, offsets_url=None, coordinate_transformations=None, coordination_values=None, **kwargs):
+    def __init__(self, img_path=None, img_url=None, img_artifact=None, offsets_path=None, offsets_url=None, offsets_artifact=None, coordinate_transformations=None, coordination_values=None, **kwargs):
         super().__init__(**kwargs)
         self._repr = make_repr(locals())
+        num_inputs = sum([1 for x in [img_path, img_url, img_artifact] if x is not None])
+        if num_inputs != 1:
+            raise ValueError(
+                "Expected one of img_path, img_url, or img_artifact to be provided")
+
+        num_inputs = sum([1 for x in [offsets_path, offsets_url, offsets_artifact] if x is not None])
+        if num_inputs > 1:
+            raise ValueError(
+                "Expected zero or one of offsets_path, offsets_url, or offsets_artifact to be provided")
+
         self._img_path = img_path
         self._img_url = img_url
+        self._img_artifact = img_artifact
         self._offsets_path = offsets_path
         self._offsets_url = offsets_url
+        self._offsets_artifact = offsets_artifact
         self._coordinate_transformations = coordinate_transformations
         self._coordination_values = coordination_values
         self.is_remote = img_url is not None
         self.local_img_uid = make_unique_filename(".ome.tif")
         self.local_offsets_uid = make_unique_filename(".offsets.json")
-        if img_url is not None and (img_path is not None or offsets_path is not None):
-            raise ValueError(
-                "Did not expect img_path or offsets_path to be provided with img_url")
+
+        if img_artifact is not None:
+            self._img_url = self.register_artifact(img_artifact)
+
+        if offsets_artifact is not None:
+            self._offsets_url = self.register_artifact(offsets_artifact)
+
 
     def convert_and_save(self, dataset_uid, obj_i, base_dir=None):
         # Only create out-directory if needed
@@ -530,22 +555,37 @@ class ObsSegmentationsOmeTiffWrapper(AbstractWrapper):
     :param \\*\\*kwargs: Keyword arguments inherited from :class:`~vitessce.wrappers.AbstractWrapper`
     """
 
-    def __init__(self, img_path=None, offsets_path=None, img_url=None, offsets_url=None, coordinate_transformations=None, obs_types_from_channel_names=None, coordination_values=None, **kwargs):
+    def __init__(self, img_path=None, img_url=None, img_artifact=None, offsets_path=None, offsets_url=None, offsets_artifact=None, coordinate_transformations=None, obs_types_from_channel_names=None, coordination_values=None, **kwargs):
         super().__init__(**kwargs)
         self._repr = make_repr(locals())
+        num_inputs = sum([1 for x in [img_path, img_url, img_artifact] if x is not None])
+        if num_inputs != 1:
+            raise ValueError(
+                "Expected one of img_path, img_url, or img_artifact to be provided")
+
+        num_inputs = sum([1 for x in [offsets_path, offsets_url, offsets_artifact] if x is not None])
+        if num_inputs > 1:
+            raise ValueError(
+                "Expected zero or one of offsets_path, offsets_url, or offsets_artifact to be provided")
+
         self._img_path = img_path
         self._img_url = img_url
+        self._img_artifact = img_artifact
         self._offsets_path = offsets_path
         self._offsets_url = offsets_url
+        self._offsets_artifact = offsets_artifact
         self._coordinate_transformations = coordinate_transformations
         self._obs_types_from_channel_names = obs_types_from_channel_names
         self._coordination_values = coordination_values
         self.is_remote = img_url is not None
         self.local_img_uid = make_unique_filename(".ome.tif")
         self.local_offsets_uid = make_unique_filename(".offsets.json")
-        if img_url is not None and (img_path is not None or offsets_path is not None):
-            raise ValueError(
-                "Did not expect img_path or offsets_path to be provided with img_url")
+
+        if img_artifact is not None:
+            self._img_url = self.register_artifact(img_artifact)
+
+        if offsets_artifact is not None:
+            self._offsets_url = self.register_artifact(offsets_artifact)
 
     def convert_and_save(self, dataset_uid, obj_i, base_dir=None):
         # Only create out-directory if needed
@@ -814,23 +854,29 @@ class ImageOmeZarrWrapper(AbstractWrapper):
     :param \\*\\*kwargs: Keyword arguments inherited from :class:`~vitessce.wrappers.AbstractWrapper`
     """
 
-    def __init__(self, img_path=None, img_url=None, coordinate_transformations=None, coordination_values=None, **kwargs):
+    def __init__(self, img_path=None, img_url=None, img_artifact=None, coordinate_transformations=None, coordination_values=None, **kwargs):
         super().__init__(**kwargs)
         self._repr = make_repr(locals())
-        if img_url is not None and img_path is not None:
-            raise ValueError(
-                "Did not expect img_path to be provided with img_url")
-        if img_url is None and img_path is None:
+
+        num_inputs = sum([1 for x in [img_path, img_url, img_artifact] if x is not None])
+        if num_inputs != 1:
             raise ValueError(
-                "Expected either img_url or img_path to be provided")
+                "Expected one of img_path, img_url, or img_artifact to be provided")
+
         self._img_path = img_path
         self._img_url = img_url
+        self._img_artifact = img_artifact
         self._coordinate_transformations = coordinate_transformations
         self._coordination_values = coordination_values
         if self._img_path is not None:
             self.is_remote = False
         else:
             self.is_remote = True
+
+        if self._img_artifact is not None:
+            # To serve as a placeholder in the config JSON URL field
+            self._img_url = self.register_artifact(img_artifact)
+
         self.local_dir_uid = make_unique_filename(".ome.zarr")
 
     def convert_and_save(self, dataset_uid, obj_i, base_dir=None):
@@ -889,24 +935,29 @@ class ObsSegmentationsOmeZarrWrapper(AbstractWrapper):
     :param \\*\\*kwargs: Keyword arguments inherited from :class:`~vitessce.wrappers.AbstractWrapper`
     """
 
-    def __init__(self, img_path=None, img_url=None, coordinate_transformations=None, coordination_values=None, obs_types_from_channel_names=None, **kwargs):
+    def __init__(self, img_path=None, img_url=None, img_artifact=None, coordinate_transformations=None, coordination_values=None, obs_types_from_channel_names=None, **kwargs):
         super().__init__(**kwargs)
         self._repr = make_repr(locals())
-        if img_url is not None and img_path is not None:
-            raise ValueError(
-                "Did not expect img_path to be provided with img_url")
-        if img_url is None and img_path is None:
+
+        num_inputs = sum([1 for x in [img_path, img_url, img_artifact] if x is not None])
+        if num_inputs != 1:
             raise ValueError(
-                "Expected either img_url or img_path to be provided")
+                "Expected one of img_path, img_url, or img_artifact to be provided")
         self._img_path = img_path
         self._img_url = img_url
+        self._img_artifact = img_artifact
         self._coordinate_transformations = coordinate_transformations
         self._obs_types_from_channel_names = obs_types_from_channel_names
         self._coordination_values = coordination_values
         if self._img_path is not None:
             self.is_remote = False
         else:
             self.is_remote = True
+
+        if self._img_artifact is not None:
+            # To serve as a placeholder in the config JSON URL field
+            self._img_url = self.register_artifact(img_artifact)
+
         self.local_dir_uid = make_unique_filename(".ome.zarr")
 
     def convert_and_save(self, dataset_uid, obj_i, base_dir=None):
@@ -956,14 +1007,16 @@ def image_file_def_creator(base_url):
 
 
 class AnnDataWrapper(AbstractWrapper):
-    def __init__(self, adata_path=None, adata_url=None, adata_store=None, ref_path=None, ref_url=None, obs_feature_matrix_path=None, feature_filter_path=None, initial_feature_filter_path=None, obs_set_paths=None, obs_set_names=None, obs_locations_path=None, obs_segmentations_path=None, obs_embedding_paths=None, obs_embedding_names=None, obs_embedding_dims=None, obs_spots_path=None, obs_points_path=None, feature_labels_path=None, obs_labels_path=None, convert_to_dense=True, coordination_values=None, obs_labels_paths=None, obs_labels_names=None, **kwargs):
+    def __init__(self, adata_path=None, adata_url=None, adata_store=None, adata_artifact=None, ref_path=None, ref_url=None, ref_artifact=None, obs_feature_matrix_path=None, feature_filter_path=None, initial_feature_filter_path=None, obs_set_paths=None, obs_set_names=None, obs_locations_path=None, obs_segmentations_path=None, obs_embedding_paths=None, obs_embedding_names=None, obs_embedding_dims=None, obs_spots_path=None, obs_points_path=None, feature_labels_path=None, obs_labels_path=None, convert_to_dense=True, coordination_values=None, obs_labels_paths=None, obs_labels_names=None, **kwargs):
         """
         Wrap an AnnData object by creating an instance of the ``AnnDataWrapper`` class.
 
         :param str adata_path: A path to an AnnData object written to a Zarr store containing single-cell experiment data.
         :param str adata_url: A remote url pointing to a zarr-backed AnnData store.
         :param adata_store: A path to pass to zarr.DirectoryStore, or an existing store instance.
         :type adata_store: str or zarr.Storage
+        :param adata_artifact: A lamindb Artifact corresponding to the AnnData object.
+        :type adata_artifact: lamindb.Artifact
         :param str obs_feature_matrix_path: Location of the expression (cell x gene) matrix, like `X` or `obsm/highly_variable_genes_subset`
         :param str feature_filter_path: A string like `var/highly_variable` used in conjunction with `obs_feature_matrix_path` if obs_feature_matrix_path points to a subset of `X` of the full `var` list.
         :param str initial_feature_filter_path: A string like `var/highly_variable` used in conjunction with `obs_feature_matrix_path` if obs_feature_matrix_path points to a subset of `X` of the full `var` list.
@@ -990,27 +1043,26 @@ def __init__(self, adata_path=None, adata_url=None, adata_store=None, ref_path=N
         self._adata_path = adata_path
         self._adata_url = adata_url
         self._adata_store = adata_store
+        self._adata_artifact = adata_artifact
 
         # For reference spec JSON with .h5ad files
         self._ref_path = ref_path
         self._ref_url = ref_url
+        self._ref_artifact = ref_artifact
 
-        if ref_path is not None or ref_url is not None:
+        if ref_path is not None or ref_url is not None or ref_artifact is not None:
             self.is_h5ad = True
         else:
             self.is_h5ad = False
 
-        if adata_store is not None and (ref_path is not None or ref_url is not None):
+        if adata_store is not None and (ref_path is not None or ref_url is not None or ref_artifact is not None):
             raise ValueError(
                 "Did not expect ref_path or ref_url to be provided with adata_store")
 
-        num_inputs = sum([1 for x in [adata_path, adata_url, adata_store] if x is not None])
-        if num_inputs > 1:
+        num_inputs = sum([1 for x in [adata_path, adata_url, adata_store, adata_artifact] if x is not None])
+        if num_inputs != 1:
             raise ValueError(
-                "Expected only one of adata_path, adata_url, or adata_store to be provided")
-        if num_inputs == 0:
-            raise ValueError(
-                "Expected one of adata_path, adata_url, or adata_store to be provided")
+                "Expected one of adata_path, adata_url, adata_artifact, or adata_store to be provided")
 
         if adata_path is not None:
             self.is_remote = False
@@ -1020,6 +1072,14 @@ def __init__(self, adata_path=None, adata_url=None, adata_store=None, ref_path=N
             self.is_remote = True
             self.is_store = False
             self.zarr_folder = None
+        elif adata_artifact is not None:
+            self.is_remote = True
+            self.is_store = False
+            self.zarr_folder = None
+            # Store artifacts on the AbstractWrapper.artifacts array for downstream access, e.g. in lamindb.save_vitessce_config
+            self._adata_url = self.register_artifact(adata_artifact)
+            if ref_artifact is not None:
+                self._ref_url = self.register_artifact(ref_artifact)
         else:
             # Store case
             self.is_remote = False