figpack 0.2.14__py3-none-any.whl → 0.2.16__py3-none-any.whl

figpack/__init__.py

@@ -2,10 +2,11 @@
 figpack - A Python package for creating shareable, interactive visualizations in the browser
 """
 
-__version__ = "0.2.14"
+__version__ = "0.2.16"
 
 from .cli import view_figure
 from .core import FigpackView, FigpackExtension, ExtensionRegistry, ExtensionView
+from .core.zarr import Group
 
 __all__ = [
     "view_figure",
@@ -13,4 +14,5 @@ __all__ = [
     "FigpackExtension",
     "ExtensionRegistry",
     "ExtensionView",
+    "Group",
 ]

figpack/core/_bundle_utils.py

@@ -7,6 +7,7 @@ import zarr
 from .figpack_view import FigpackView
 from .figpack_extension import ExtensionRegistry
 from .extension_view import ExtensionView
+from .zarr import Group, _check_zarr_version
 
 thisdir = pathlib.Path(__file__).parent.resolve()
 
@@ -45,24 +46,32 @@ def prepare_figure_bundle(
                 target_sub = target / subitem.name
                 target_sub.write_bytes(subitem.read_bytes())
 
-    # Write the view data to the Zarr group
-    zarr_group = zarr.open_group(
-        pathlib.Path(tmpdir) / "data.zarr",
-        mode="w",
-        synchronizer=zarr.ThreadSynchronizer(),
-    )
-    view._write_to_zarr_group(zarr_group)
+    # If we are using zarr 3, then we set the default zarr format to 2 temporarily
+    # because we only support version 2 on the frontend right now.
 
-    # Add title and description as attributes on the top-level zarr group
-    zarr_group.attrs["title"] = title
-    if description is not None:
-        zarr_group.attrs["description"] = description
+    if _check_zarr_version() == 3:
+        old_default_zarr_format = zarr.config.get("default_zarr_format")
+        zarr.config.set({"default_zarr_format": 2})
 
-    # Discover and write extension JavaScript files
-    required_extensions = _discover_required_extensions(view)
-    _write_extension_files(required_extensions, tmpdir)
+    try:
+        # Write the view data to the Zarr group
+        zarr_group = zarr.open_group(pathlib.Path(tmpdir) / "data.zarr", mode="w")
+        zarr_group = Group(zarr_group)
+        view._write_to_zarr_group(zarr_group)
 
-    zarr.consolidate_metadata(zarr_group.store)
+        # Add title and description as attributes on the top-level zarr group
+        zarr_group.attrs["title"] = title
+        if description is not None:
+            zarr_group.attrs["description"] = description
+
+        # Discover and write extension JavaScript files
+        required_extensions = _discover_required_extensions(view)
+        _write_extension_files(required_extensions, tmpdir)
+
+        zarr.consolidate_metadata(zarr_group._zarr_group.store)
+    finally:
+        if _check_zarr_version() == 3:
+            zarr.config.set({"default_zarr_format": old_default_zarr_format})
 
 
 def _discover_required_extensions(view: FigpackView) -> Set[str]:

figpack/core/_view_figure.py

@@ -109,7 +109,7 @@ def view_figure(figure_path: str, port: Union[int, None] = None) -> None:
 
         try:
             with tarfile.open(figure_path, "r:gz") as tar:
-                tar.extractall(temp_path)
+                tar.extractall(temp_path, filter="data")
 
             # Count extracted files
             extracted_files = list(temp_path.rglob("*"))

figpack/core/extension_view.py

@@ -2,9 +2,9 @@
 Base class for views that use figpack extensions
 """
 
-import zarr
 from .figpack_view import FigpackView
 from .figpack_extension import ExtensionRegistry
+from ..core.zarr import Group
 
 
 class ExtensionView(FigpackView):
@@ -32,7 +32,7 @@ class ExtensionView(FigpackView):
             )
         self.extension = extension
 
-    def _write_to_zarr_group(self, group: zarr.Group) -> None:
+    def _write_to_zarr_group(self, group: Group) -> None:
         """
         Write the extension view metadata to a Zarr group.
         Subclasses should call super()._write_to_zarr_group(group) first,

figpack/core/figpack_view.py

@@ -7,7 +7,7 @@ import random
 import string
 from typing import Union
 
-import zarr
+from .zarr import Group
 
 
 class FigpackView:
@@ -155,7 +155,7 @@ class FigpackView:
 
         _save_figure(self, output_path, title=title)
 
-    def _write_to_zarr_group(self, group: zarr.Group) -> None:
+    def _write_to_zarr_group(self, group: Group) -> None:
         """
         Write the view data to a Zarr group. Must be implemented by subclasses.
 

figpack/core/zarr.py

@@ -0,0 +1,61 @@
+from typing import Any, Dict
+import zarr
+
+
+_UNSPECIFIED = object()
+
+
+class Group:
+    def __init__(self, zarr_group: zarr.Group):
+        self._zarr_group = zarr_group
+
+    def create_group(self, name: str) -> "Group":
+        return Group(self._zarr_group.create_group(name))
+
+    def create_dataset(
+        self,
+        name: str,
+        *,
+        data=_UNSPECIFIED,
+        dtype=_UNSPECIFIED,
+        chunks=_UNSPECIFIED,
+        compressor=_UNSPECIFIED,
+    ) -> None:
+        kwargs = {}
+        if data is not _UNSPECIFIED:
+            kwargs["data"] = data
+        if dtype is not _UNSPECIFIED:
+            kwargs["dtype"] = dtype
+        if chunks is not _UNSPECIFIED:
+            kwargs["chunks"] = chunks
+        if compressor is not _UNSPECIFIED:
+            kwargs["compressor"] = compressor
+        if _check_zarr_version() == 2:
+            self._zarr_group.create_dataset(name, **kwargs)
+        elif _check_zarr_version() == 3:
+            self._zarr_group.create_array(name, **kwargs)
+        else:
+            raise RuntimeError("Unsupported Zarr version")
+
+    @property
+    def attrs(self) -> Dict[str, Any]:
+        return self._zarr_group.attrs
+
+    def __getitem__(self, key: str) -> Any:
+        return self._zarr_group[key]
+
+    # implement in operator
+    def __contains__(self, key: str) -> bool:
+        return key in self._zarr_group
+
+    def __iter__(self):
+        return iter(self._zarr_group)
+
+    def __reversed__(self):
+        return reversed(self._zarr_group)
+
+
+def _check_zarr_version():
+    version = zarr.__version__
+    major_version = int(version.split(".")[0])
+    return major_version