flytekitplugins-xarray-zarr 1.16.9__py3-none-any.whl

flytekitplugins/xarray/__init__.py

@@ -0,0 +1,14 @@
+"""
+.. currentmodule:: flytekitplugins.xarray
+
+This package contains things that are useful when extending Flytekit.
+
+.. autosummary::
+   :template: custom.rst
+   :toctree: generated/
+
+   XarrayDaZarrTypeTransformer
+   XarrayZarrTypeTransformer
+"""
+
+from .xarray_transformers import XarrayDaZarrTypeTransformer, XarrayZarrTypeTransformer

flytekitplugins/xarray/xarray_transformers.py

@@ -0,0 +1,102 @@
+import typing
+
+import dask.distributed as dd
+
+import xarray as xr
+from flytekit import (
+    Blob,
+    BlobMetadata,
+    BlobType,
+    FlyteContext,
+    Literal,
+    LiteralType,
+    Scalar,
+)
+from flytekit.extend import TypeEngine, TypeTransformer
+
+
+class XarrayZarrTypeTransformer(TypeTransformer[xr.Dataset]):
+    _TYPE_INFO = BlobType(format="binary", dimensionality=BlobType.BlobDimensionality.MULTIPART)
+
+    def __init__(self) -> None:
+        super().__init__(name="Xarray Dataset", t=xr.Dataset)
+
+    def get_literal_type(self, t: typing.Type[xr.Dataset]) -> LiteralType:
+        return LiteralType(blob=self._TYPE_INFO)
+
+    def to_literal(
+        self,
+        ctx: FlyteContext,
+        python_val: xr.Dataset,
+        python_type: typing.Type[xr.Dataset],
+        expected: LiteralType,
+    ) -> Literal:
+        remote_dir = ctx.file_access.get_random_remote_path("data.zarr")
+        # Opening with the dask client will attach the client eliminating the
+        # need for users to connect to the client if a task tasks a xr.Dataset
+        # type.
+        with dd.Client(timeout=120):
+            python_val.to_zarr(remote_dir, mode="w")
+        return Literal(scalar=Scalar(blob=Blob(uri=remote_dir, metadata=BlobMetadata(type=self._TYPE_INFO))))
+
+    def to_python_value(
+        self,
+        ctx: FlyteContext,
+        lv: Literal,
+        expected_python_type: typing.Type[xr.Dataset],
+    ) -> xr.Dataset:
+        return xr.open_zarr(lv.scalar.blob.uri)
+
+    def to_html(
+        self,
+        ctx: FlyteContext,
+        python_val: xr.Dataset,
+        expected_python_type: LiteralType,
+    ) -> str:
+        return python_val._repr_html_()
+
+
+class XarrayDaZarrTypeTransformer(TypeTransformer[xr.DataArray]):
+    _TYPE_INFO = BlobType(format="binary", dimensionality=BlobType.BlobDimensionality.MULTIPART)
+
+    def __init__(self) -> None:
+        super().__init__(name="Xarray DataArray", t=xr.DataArray)
+
+    def get_literal_type(self, t: typing.Type[xr.DataArray]) -> LiteralType:
+        return LiteralType(blob=self._TYPE_INFO)
+
+    def to_literal(
+        self,
+        ctx: FlyteContext,
+        python_val: xr.DataArray,
+        python_type: typing.Type[xr.DataArray],
+        expected: LiteralType,
+    ) -> Literal:
+        remote_dir = ctx.file_access.get_random_remote_path("data.zarr")
+        # Opening with the dask client will attach the client eliminating the
+        # need for users to connect to the client if a task tasks a xr.Dataset
+        # type.
+        with dd.Client(timeout=120):
+            python_val.to_zarr(remote_dir, mode="w")
+        return Literal(scalar=Scalar(blob=Blob(uri=remote_dir, metadata=BlobMetadata(type=self._TYPE_INFO))))
+
+    def to_python_value(
+        self,
+        ctx: FlyteContext,
+        lv: Literal,
+        expected_python_type: typing.Type[xr.DataArray],
+    ) -> xr.DataArray:
+        # xr.open_zarr always opens a dataset, so we take the first variable
+        return list(xr.open_zarr(lv.scalar.blob.uri).data_vars.values())[0]
+
+    def to_html(
+        self,
+        ctx: FlyteContext,
+        python_val: xr.DataArray,
+        expected_python_type: LiteralType,
+    ) -> str:
+        return python_val._repr_html_()
+
+
+TypeEngine.register(XarrayZarrTypeTransformer())
+TypeEngine.register(XarrayDaZarrTypeTransformer())

flytekitplugins_xarray_zarr-1.16.9-py3.14-nspkg.pth

@@ -0,0 +1 @@
+import sys, types, os;p = os.path.join(sys._getframe(1).f_locals['sitedir'], *('flytekitplugins',));importlib = __import__('importlib.util');__import__('importlib.machinery');m = sys.modules.setdefault('flytekitplugins', importlib.util.module_from_spec(importlib.machinery.PathFinder.find_spec('flytekitplugins', [os.path.dirname(p)])));m = m or sys.modules.setdefault('flytekitplugins', types.ModuleType('flytekitplugins'));mp = (m or []) and m.__dict__.setdefault('__path__',[]);(p not in mp) and mp.append(p)

flytekitplugins_xarray_zarr-1.16.9.dist-info/METADATA

@@ -0,0 +1,29 @@
+Metadata-Version: 2.4
+Name: flytekitplugins-xarray-zarr
+Version: 1.16.9
+Summary: Xarray Zarr plugin for flytekit
+Author: flyteorg
+Author-email: admin@flyte.org
+License: apache2
+Classifier: Intended Audience :: Science/Research
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Topic :: Scientific/Engineering
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Software Development
+Classifier: Topic :: Software Development :: Libraries
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Requires-Python: >=3.9
+Requires-Dist: dask[distributed]>=2022.10.2
+Requires-Dist: flytekit<2.0.0,>=1.3.0b2
+Requires-Dist: xarray
+Requires-Dist: zarr
+Dynamic: author
+Dynamic: author-email
+Dynamic: classifier
+Dynamic: license
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary

flytekitplugins_xarray_zarr-1.16.9.dist-info/RECORD

@@ -0,0 +1,9 @@
+flytekitplugins_xarray_zarr-1.16.9-py3.14-nspkg.pth,sha256=ya_zoT0ZM023a2lpUsRaqXsgDEknLucLLOWbbtFBPw4,512
+flytekitplugins/xarray/__init__.py,sha256=p9pV5dDng3ujPCPISggi-YWNb2gJRWfeM_t0Onum_hA,352
+flytekitplugins/xarray/xarray_transformers.py,sha256=qi04uyukJ6EMZWoBi8A_uqQsFF3izZCAHxvGEkq7eOQ,3308
+flytekitplugins_xarray_zarr-1.16.9.dist-info/METADATA,sha256=YvXntLo4QfyZJXaKgEgVPX0anycAx44de7rTZPZlCOo,1009
+flytekitplugins_xarray_zarr-1.16.9.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+flytekitplugins_xarray_zarr-1.16.9.dist-info/entry_points.txt,sha256=JJHdCPzY51JrqrBaJJYovbBKXeRYCWHIF41dtfN5G8U,51
+flytekitplugins_xarray_zarr-1.16.9.dist-info/namespace_packages.txt,sha256=b2gTzXFHuV5afchaJ28HueRhYzj-8x5Ytr8n2twnAik,16
+flytekitplugins_xarray_zarr-1.16.9.dist-info/top_level.txt,sha256=b2gTzXFHuV5afchaJ28HueRhYzj-8x5Ytr8n2twnAik,16
+flytekitplugins_xarray_zarr-1.16.9.dist-info/RECORD,,

flytekitplugins_xarray_zarr-1.16.9.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (80.9.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

flytekitplugins_xarray_zarr-1.16.9.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[flytekit.plugins]
+xarray = flytekitplugins.xarray

flytekitplugins_xarray_zarr-1.16.9.dist-info/namespace_packages.txt

@@ -0,0 +1 @@
+flytekitplugins

flytekitplugins_xarray_zarr-1.16.9.dist-info/top_level.txt

@@ -0,0 +1 @@
+flytekitplugins