Use anywidget-based IPC for zarr gets which do not require serving da…

…ta on localhost (#329)

* WIP: zarr get using anywidget.experimental command/invoke works

* Unused zarr js import

* Support local store

* Lint

* Cleanup

* Lint

* Bump version

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "vitessce"
-version = "3.2.5"
+version = "3.2.6"
 authors = [
   { name="Mark Keller", email="mark_keller@hms.harvard.edu" },
 ]
@@ -73,7 +73,7 @@ docs = [
 ]
 all = [
   'jupyter-server-proxy>=1.5.2',
-  'anywidget>=0.9.3',
+  'anywidget>=0.9.10',
   'uvicorn>=0.17.0',
   'ujson>=4.0.1',
   'starlette==0.14.0',

vitessce/config.py

@@ -227,6 +227,16 @@ def get_routes(self):
 
         return routes
 
+    def get_stores(self, base_url=None):
+        stores = {}
+        for obj in self.objs:
+            stores = {
+                **stores,
+                **obj.get_stores(base_url)
+            }
+
+        return stores
+
 
 class VitessceConfigViewHConcat:
     """
@@ -1489,6 +1499,21 @@ def get_routes(self):
             routes += d.get_routes()
         return routes
 
+    def get_stores(self, base_url=None):
+        """
+        Convert the routes for this view config from the datasets.
+
+        :returns: A list of server routes.
+        :rtype: list[starlette.routing.Route]
+        """
+        stores = {}
+        for d in self.config["datasets"]:
+            stores = {
+                **stores,
+                **d.get_stores(base_url)
+            }
+        return stores
+
     def to_python(self):
         """
         Convert the VitessceConfig instance to a one-line Python code snippet that can be used to generate it.

vitessce/widget.py

@@ -4,7 +4,7 @@
 
 # Widget dependencies
 import anywidget
-from traitlets import Unicode, Dict, Int, Bool
+from traitlets import Unicode, Dict, List, Int, Bool
 import time
 import uuid
 
@@ -201,6 +201,7 @@ def get_uid_str(uid):
     const customJsUrl = view.model.get('custom_js_url');
     const pluginEsm = view.model.get('plugin_esm');
     const remountOnUidChange = view.model.get('remount_on_uid_change');
+    const storeUrls = view.model.get('store_urls');
 
     const pkgName = (jsDevMode ? "@vitessce/dev" : "vitessce");
 
@@ -224,6 +225,19 @@ def get_uid_str(uid):
     let pluginFileTypes;
     let pluginJointFileTypes;
 
+    const stores = Object.fromEntries(
+        storeUrls.map(storeUrl => ([
+            storeUrl,
+            {
+                async get(key) {
+                    const [data, buffers] = await view.experimental.invoke("_zarr_get", [storeUrl, key]);
+                    if (!data.success) return undefined;
+                    return buffers[0].buffer;
+                },
+            }
+        ])),
+    );
+
     try {
         const pluginEsmUrl = URL.createObjectURL(new Blob([pluginEsm], { type: "text/javascript" }));
         const pluginModule = (await import(pluginEsmUrl)).default;
@@ -310,7 +324,7 @@ def get_uid_str(uid):
         const vitessceProps = {
             height, theme, config, onConfigChange, validateConfig,
             pluginViewTypes, pluginCoordinationTypes, pluginFileTypes, pluginJointFileTypes,
-            remountOnUidChange,
+            remountOnUidChange, stores,
         };
 
         return e('div', { ref: divRef, style: { height: height + 'px' } },
@@ -383,13 +397,15 @@ class VitessceWidget(anywidget.AnyWidget):
 
     next_port = DEFAULT_PORT
 
-    js_package_version = Unicode('3.3.7').tag(sync=True)
+    js_package_version = Unicode('3.3.12').tag(sync=True)
     js_dev_mode = Bool(False).tag(sync=True)
     custom_js_url = Unicode('').tag(sync=True)
     plugin_esm = Unicode(DEFAULT_PLUGIN_ESM).tag(sync=True)
     remount_on_uid_change = Bool(True).tag(sync=True)
 
-    def __init__(self, config, height=600, theme='auto', uid=None, port=None, proxy=False, js_package_version='3.3.7', js_dev_mode=False, custom_js_url='', plugin_esm=DEFAULT_PLUGIN_ESM, remount_on_uid_change=True):
+    store_urls = List(trait=Unicode(''), default_value=[]).tag(sync=True)
+
+    def __init__(self, config, height=600, theme='auto', uid=None, port=None, proxy=False, js_package_version='3.3.12', js_dev_mode=False, custom_js_url='', plugin_esm=DEFAULT_PLUGIN_ESM, remount_on_uid_change=True):
         """
         Construct a new Vitessce widget.
 
@@ -422,13 +438,15 @@ def __init__(self, config, height=600, theme='auto', uid=None, port=None, proxy=
         config_dict = config.to_dict(base_url=base_url)
         routes = config.get_routes()
 
+        self._stores = config.get_stores(base_url=base_url)
+
         uid_str = get_uid_str(uid)
 
         super(VitessceWidget, self).__init__(
             config=config_dict, height=height, theme=theme, proxy=proxy,
             js_package_version=js_package_version, js_dev_mode=js_dev_mode, custom_js_url=custom_js_url,
             plugin_esm=plugin_esm, remount_on_uid_change=remount_on_uid_change,
-            uid=uid_str,
+            uid=uid_str, store_urls=list(self._stores.keys())
         )
 
         serve_routes(config, routes, use_port)
@@ -460,10 +478,20 @@ def close(self):
         self.config_obj.stop_server(self.port)
         super().close()
 
+    @anywidget.experimental.command
+    def _zarr_get(self, params, buffers):
+        [store_url, key] = params
+        store = self._stores[store_url]
+        try:
+            buffers = [store[key.lstrip("/")]]
+        except KeyError:
+            buffers = []
+        return {"success": len(buffers) == 1}, buffers
+
 # Launch Vitessce using plain HTML representation (no ipywidgets)
 
 
-def ipython_display(config, height=600, theme='auto', base_url=None, host_name=None, uid=None, port=None, proxy=False, js_package_version='3.3.7', js_dev_mode=False, custom_js_url='', plugin_esm=DEFAULT_PLUGIN_ESM, remount_on_uid_change=True):
+def ipython_display(config, height=600, theme='auto', base_url=None, host_name=None, uid=None, port=None, proxy=False, js_package_version='3.3.12', js_dev_mode=False, custom_js_url='', plugin_esm=DEFAULT_PLUGIN_ESM, remount_on_uid_change=True):
     from IPython.display import display, HTML
     uid_str = "vitessce" + get_uid_str(uid)
 
@@ -485,6 +513,7 @@ def ipython_display(config, height=600, theme='auto', base_url=None, host_name=N
         "height": height,
         "theme": theme,
         "config": config_dict,
+        "store_urls": [],
     }
 
     # We need to clean up the React and DOM state in any case in which

vitessce/wrappers.py

@@ -3,6 +3,7 @@
 import tempfile
 from uuid import uuid4
 from pathlib import PurePath, PurePosixPath
+import zarr
 
 from .constants import (
     norm_enum,
@@ -41,9 +42,11 @@ def __init__(self, **kwargs):
         self.out_dir = kwargs['out_dir'] if 'out_dir' in kwargs else tempfile.mkdtemp(
         )
         self.routes = []
-        self.is_remote = False
+        self.is_remote = False  # TODO: change to needs_localhost_serving for clarity
+        self.is_store = False  # TODO: change to needs_store_registration for clarity
         self.file_def_creators = []
         self.base_dir = None
+        self.stores = {}
         self._request_init = kwargs['request_init'] if 'request_init' in kwargs else None
 
     def __repr__(self):
@@ -71,6 +74,20 @@ def get_routes(self):
         """
         return self.routes
 
+    def get_stores(self, base_url):
+        """
+        Obtain the stores that have been created for this wrapper class.
+
+        :returns: A dictionary that maps file URLs to Zarr Store objects.
+        :rtype: dict[str, zarr.Store]
+        """
+        relative_stores = self.stores
+        absolute_stores = {}
+        for relative_url, store in relative_stores.items():
+            absolute_url = base_url + relative_url
+            absolute_stores[absolute_url] = store
+        return absolute_stores
+
     def get_file_defs(self, base_url):
         """
         Obtain the file definitions for this wrapper class.
@@ -111,6 +128,30 @@ def get_local_dir_url(self, base_url, dataset_uid, obj_i, local_dir_path, local_
             return self._get_url_simple(base_url, file_path_to_url_path(local_dir_path, prepend_slash=False))
         return self._get_url(base_url, dataset_uid, obj_i, local_dir_uid)
 
+    def register_zarr_store(self, dataset_uid, obj_i, store_or_local_dir_path, local_dir_uid):
+        if not self.is_remote and self.is_store:
+            # Set up `store` and `local_dir_path` variables.
+            if isinstance(store_or_local_dir_path, str):
+                # TODO: use zarr.FSStore if fsspec is installed?
+                store = zarr.DirectoryStore(store_or_local_dir_path)
+                local_dir_path = store_or_local_dir_path
+            else:
+                # TODO: check that store_or_local_dir_path is a zarr.Store or StoreLike?
+                store = store_or_local_dir_path
+                # A store instance was passed directly, so there is no local directory path.
+                # Instead we just make one up using _get_route_str but it could be any string.
+                local_dir_path = self._get_route_str(dataset_uid, obj_i, local_dir_uid)
+
+            # Register the store on the same route path
+            # that will be used for the "url" field in the file definition.
+            if self.base_dir is None:
+                route_path = self._get_route_str(dataset_uid, obj_i, local_dir_uid)
+            else:
+                route_path = file_path_to_url_path(local_dir_path)
+                local_dir_path = join(self.base_dir, local_dir_path)
+
+            self.stores[route_path] = store
+
     def get_local_dir_route(self, dataset_uid, obj_i, local_dir_path, local_dir_uid):
         """
         Obtain the Mount for some local directory
@@ -896,12 +937,14 @@ def image_file_def_creator(base_url):
 
 
 class AnnDataWrapper(AbstractWrapper):
-    def __init__(self, adata_path=None, adata_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, 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.FSStore, or an existing store instance.
+        :type adata_store: str or zarr.Storage
         :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.
@@ -927,18 +970,30 @@ def __init__(self, adata_path=None, adata_url=None, obs_feature_matrix_path=None
         self._repr = make_repr(locals())
         self._adata_path = adata_path
         self._adata_url = adata_url
-        if adata_url is not None and (adata_path is not None):
+        self._adata_store = adata_store
+
+        num_inputs = sum([1 for x in [adata_path, adata_url, adata_store] if x is not None])
+        if num_inputs > 1:
             raise ValueError(
-                "Did not expect adata_url to be provided with adata_path")
-        if adata_url is None and (adata_path is None):
+                "Expected only one of adata_path, adata_url, or adata_store to be provided")
+        if num_inputs == 0:
             raise ValueError(
-                "Expected either adata_url or adata_path to be provided")
+                "Expected one of adata_path, adata_url, or adata_store to be provided")
+
         if adata_path is not None:
             self.is_remote = False
+            self.is_store = False
             self.zarr_folder = 'anndata.zarr'
-        else:
+        elif adata_url is not None:
             self.is_remote = True
+            self.is_store = False
             self.zarr_folder = None
+        else:
+            # Store case
+            self.is_remote = False
+            self.is_store = True
+            self.zarr_folder = None
+
         self.local_dir_uid = make_unique_filename(".adata.zarr")
         self._expression_matrix = obs_feature_matrix_path
         self._cell_set_obs_names = obs_set_names
@@ -978,6 +1033,9 @@ def convert_and_save(self, dataset_uid, obj_i, base_dir=None):
     def make_anndata_routes(self, dataset_uid, obj_i):
         if self.is_remote:
             return []
+        elif self.is_store:
+            self.register_zarr_store(dataset_uid, obj_i, self._adata_store, self.local_dir_uid)
+            return []
         else:
             return self.get_local_dir_route(dataset_uid, obj_i, self._adata_path, self.local_dir_uid)