Merge pull request #560 from bioimage-io/zip2bytes

Add function to write models directly to streams

bioimageio/spec/__init__.py

@@ -25,6 +25,9 @@
 from ._package import (
     save_bioimageio_package_as_folder as save_bioimageio_package_as_folder,
 )
+from ._package import (
+    save_bioimageio_package_to_stream as save_bioimageio_package_to_stream,
+)
 from .application import AnyApplicationDescr as AnyApplicationDescr
 from .application import ApplicationDescr as ApplicationDescr
 from .collection import AnyCollectionDescr as AnyCollectionDescr

bioimageio/spec/_internal/io_utils.py

@@ -3,7 +3,7 @@
 import warnings
 from contextlib import nullcontext
 from pathlib import Path
-from typing import Any, Dict, Iterable, Mapping, Optional, TextIO, Union, cast
+from typing import IO, Any, Dict, Iterable, Mapping, Optional, TextIO, Union, cast
 from zipfile import ZipFile, is_zipfile
 
 import numpy
@@ -170,7 +170,7 @@ def unzip(
 
 
 def write_zip(
-    path: FilePath,
+    path: Union[FilePath, IO[bytes]],
     content: Mapping[FileName, Union[str, FilePath, Dict[Any, Any]]],
     *,
     compression: int,

bioimageio/spec/_package.py

@@ -1,9 +1,10 @@
 import collections.abc
+from io import BytesIO
 import re
 import shutil
 from pathlib import Path
 from tempfile import NamedTemporaryFile, mkdtemp
-from typing import Dict, Literal, Optional, Sequence, Union, cast
+from typing import IO, Dict, Literal, Optional, Sequence, Union, cast
 from zipfile import ZIP_DEFLATED
 
 from pydantic import DirectoryPath, FilePath, NewPath
@@ -257,3 +258,57 @@ def save_bioimageio_package(
         )
 
     return output_path
+
+
+def save_bioimageio_package_to_stream(
+    source: Union[BioimageioYamlSource, ResourceDescr],
+    /,
+    *,
+    compression: int = ZIP_DEFLATED,
+    compression_level: int = 1,
+    output_stream: Union[IO[bytes], None] = None,
+    weights_priority_order: Optional[  # model only
+        Sequence[
+            Literal[
+                "keras_hdf5",
+                "onnx",
+                "pytorch_state_dict",
+                "tensorflow_js",
+                "tensorflow_saved_model_bundle",
+                "torchscript",
+            ]
+        ]
+    ] = None,
+) -> IO[bytes]:
+    """Package a bioimageio resource into a stream.
+
+    Args:
+        rd: bioimageio resource description
+        compression: The numeric constant of compression method.
+        compression_level: Compression level to use when writing files to the archive.
+                           See https://docs.python.org/3/library/zipfile.html#zipfile.ZipFile
+        output_stream: stream to write package to
+        weights_priority_order: If given only the first weights format present in the model is included.
+                                If none of the prioritized weights formats is found all are included.
+
+    Note: this function bypasses safety checks and does not load/validate the model after writing.
+
+    Returns:
+        stream of zipped bioimageio package
+    """
+    if output_stream is None:
+        output_stream = BytesIO()
+
+    package_content = _prepare_resource_package(
+        source,
+        weights_priority_order=weights_priority_order,
+    )
+
+    write_zip(
+        output_stream,
+        package_content,
+        compression=compression,
+        compression_level=compression_level,
+    )
+
+    return output_stream

bioimageio/spec/application/__init__.py

@@ -1,20 +1,20 @@
-# autogen: start
-"""
-implementaions of all released minor versions are available in submodules:
-- application v0_2: `bioimageio.spec.application.v0_2.ApplicationDescr` [user documentation](../../../user_docs/application_descr_v0-2.md)
-- application v0_3: `bioimageio.spec.application.v0_3.ApplicationDescr` [user documentation](../../../user_docs/application_descr_v0-3.md)
-"""
-from typing import Union
-
-from pydantic import Discriminator
-from typing_extensions import Annotated
-
-from .v0_2 import ApplicationDescr as ApplicationDescr_v0_2
-from .v0_3 import ApplicationDescr as ApplicationDescr
-from .v0_3 import ApplicationDescr as ApplicationDescr_v0_3
-
-AnyApplicationDescr = Annotated[
-    Union[ApplicationDescr_v0_2, ApplicationDescr_v0_3], Discriminator("format_version")
-]
-"""Union of any released application desription"""
-# autogen: stop
+# autogen: start
+"""
+implementaions of all released minor versions are available in submodules:
+- application v0_2: `bioimageio.spec.application.v0_2.ApplicationDescr` [user documentation](../../../user_docs/application_descr_v0-2.md)
+- application v0_3: `bioimageio.spec.application.v0_3.ApplicationDescr` [user documentation](../../../user_docs/application_descr_v0-3.md)
+"""
+from typing import Union
+
+from pydantic import Discriminator
+from typing_extensions import Annotated
+
+from .v0_2 import ApplicationDescr as ApplicationDescr_v0_2
+from .v0_3 import ApplicationDescr as ApplicationDescr
+from .v0_3 import ApplicationDescr as ApplicationDescr_v0_3
+
+AnyApplicationDescr = Annotated[
+    Union[ApplicationDescr_v0_2, ApplicationDescr_v0_3], Discriminator("format_version")
+]
+"""Union of any released application desription"""
+# autogen: stop

bioimageio/spec/collection/__init__.py

@@ -1,20 +1,20 @@
-# autogen: start
-"""
-implementaions of all released minor versions are available in submodules:
-- collection v0_2: `bioimageio.spec.collection.v0_2.CollectionDescr` [user documentation](../../../user_docs/collection_descr_v0-2.md)
-- collection v0_3: `bioimageio.spec.collection.v0_3.CollectionDescr` [user documentation](../../../user_docs/collection_descr_v0-3.md)
-"""
-from typing import Union
-
-from pydantic import Discriminator
-from typing_extensions import Annotated
-
-from .v0_2 import CollectionDescr as CollectionDescr_v0_2
-from .v0_3 import CollectionDescr as CollectionDescr
-from .v0_3 import CollectionDescr as CollectionDescr_v0_3
-
-AnyCollectionDescr = Annotated[
-    Union[CollectionDescr_v0_2, CollectionDescr_v0_3], Discriminator("format_version")
-]
-"""Union of any released collection desription"""
-# autogen: stop
+# autogen: start
+"""
+implementaions of all released minor versions are available in submodules:
+- collection v0_2: `bioimageio.spec.collection.v0_2.CollectionDescr` [user documentation](../../../user_docs/collection_descr_v0-2.md)
+- collection v0_3: `bioimageio.spec.collection.v0_3.CollectionDescr` [user documentation](../../../user_docs/collection_descr_v0-3.md)
+"""
+from typing import Union
+
+from pydantic import Discriminator
+from typing_extensions import Annotated
+
+from .v0_2 import CollectionDescr as CollectionDescr_v0_2
+from .v0_3 import CollectionDescr as CollectionDescr
+from .v0_3 import CollectionDescr as CollectionDescr_v0_3
+
+AnyCollectionDescr = Annotated[
+    Union[CollectionDescr_v0_2, CollectionDescr_v0_3], Discriminator("format_version")
+]
+"""Union of any released collection desription"""
+# autogen: stop

bioimageio/spec/dataset/__init__.py

@@ -1,20 +1,20 @@
-# autogen: start
-"""
-implementaions of all released minor versions are available in submodules:
-- dataset v0_2: `bioimageio.spec.dataset.v0_2.DatasetDescr` [user documentation](../../../user_docs/dataset_descr_v0-2.md)
-- dataset v0_3: `bioimageio.spec.dataset.v0_3.DatasetDescr` [user documentation](../../../user_docs/dataset_descr_v0-3.md)
-"""
-from typing import Union
-
-from pydantic import Discriminator
-from typing_extensions import Annotated
-
-from .v0_2 import DatasetDescr as DatasetDescr_v0_2
-from .v0_3 import DatasetDescr as DatasetDescr
-from .v0_3 import DatasetDescr as DatasetDescr_v0_3
-
-AnyDatasetDescr = Annotated[
-    Union[DatasetDescr_v0_2, DatasetDescr_v0_3], Discriminator("format_version")
-]
-"""Union of any released dataset desription"""
-# autogen: stop
+# autogen: start
+"""
+implementaions of all released minor versions are available in submodules:
+- dataset v0_2: `bioimageio.spec.dataset.v0_2.DatasetDescr` [user documentation](../../../user_docs/dataset_descr_v0-2.md)
+- dataset v0_3: `bioimageio.spec.dataset.v0_3.DatasetDescr` [user documentation](../../../user_docs/dataset_descr_v0-3.md)
+"""
+from typing import Union
+
+from pydantic import Discriminator
+from typing_extensions import Annotated
+
+from .v0_2 import DatasetDescr as DatasetDescr_v0_2
+from .v0_3 import DatasetDescr as DatasetDescr
+from .v0_3 import DatasetDescr as DatasetDescr_v0_3
+
+AnyDatasetDescr = Annotated[
+    Union[DatasetDescr_v0_2, DatasetDescr_v0_3], Discriminator("format_version")
+]
+"""Union of any released dataset desription"""
+# autogen: stop

bioimageio/spec/generic/__init__.py

@@ -1,20 +1,20 @@
-# autogen: start
-"""
-implementaions of all released minor versions are available in submodules:
-- generic v0_2: `bioimageio.spec.generic.v0_2.GenericDescr` [user documentation](../../../user_docs/generic_descr_v0-2.md)
-- generic v0_3: `bioimageio.spec.generic.v0_3.GenericDescr` [user documentation](../../../user_docs/generic_descr_v0-3.md)
-"""
-from typing import Union
-
-from pydantic import Discriminator
-from typing_extensions import Annotated
-
-from .v0_2 import GenericDescr as GenericDescr_v0_2
-from .v0_3 import GenericDescr as GenericDescr
-from .v0_3 import GenericDescr as GenericDescr_v0_3
-
-AnyGenericDescr = Annotated[
-    Union[GenericDescr_v0_2, GenericDescr_v0_3], Discriminator("format_version")
-]
-"""Union of any released generic desription"""
-# autogen: stop
+# autogen: start
+"""
+implementaions of all released minor versions are available in submodules:
+- generic v0_2: `bioimageio.spec.generic.v0_2.GenericDescr` [user documentation](../../../user_docs/generic_descr_v0-2.md)
+- generic v0_3: `bioimageio.spec.generic.v0_3.GenericDescr` [user documentation](../../../user_docs/generic_descr_v0-3.md)
+"""
+from typing import Union
+
+from pydantic import Discriminator
+from typing_extensions import Annotated
+
+from .v0_2 import GenericDescr as GenericDescr_v0_2
+from .v0_3 import GenericDescr as GenericDescr
+from .v0_3 import GenericDescr as GenericDescr_v0_3
+
+AnyGenericDescr = Annotated[
+    Union[GenericDescr_v0_2, GenericDescr_v0_3], Discriminator("format_version")
+]
+"""Union of any released generic desription"""
+# autogen: stop

bioimageio/spec/model/__init__.py

@@ -1,20 +1,20 @@
-# autogen: start
-"""
-implementaions of all released minor versions are available in submodules:
-- model v0_4: `bioimageio.spec.model.v0_4.ModelDescr` [user documentation](../../../user_docs/model_descr_v0-4.md)
-- model v0_5: `bioimageio.spec.model.v0_5.ModelDescr` [user documentation](../../../user_docs/model_descr_v0-5.md)
-"""
-from typing import Union
-
-from pydantic import Discriminator
-from typing_extensions import Annotated
-
-from .v0_4 import ModelDescr as ModelDescr_v0_4
-from .v0_5 import ModelDescr as ModelDescr
-from .v0_5 import ModelDescr as ModelDescr_v0_5
-
-AnyModelDescr = Annotated[
-    Union[ModelDescr_v0_4, ModelDescr_v0_5], Discriminator("format_version")
-]
-"""Union of any released model desription"""
-# autogen: stop
+# autogen: start
+"""
+implementaions of all released minor versions are available in submodules:
+- model v0_4: `bioimageio.spec.model.v0_4.ModelDescr` [user documentation](../../../user_docs/model_descr_v0-4.md)
+- model v0_5: `bioimageio.spec.model.v0_5.ModelDescr` [user documentation](../../../user_docs/model_descr_v0-5.md)
+"""
+from typing import Union
+
+from pydantic import Discriminator
+from typing_extensions import Annotated
+
+from .v0_4 import ModelDescr as ModelDescr_v0_4
+from .v0_5 import ModelDescr as ModelDescr
+from .v0_5 import ModelDescr as ModelDescr_v0_5
+
+AnyModelDescr = Annotated[
+    Union[ModelDescr_v0_4, ModelDescr_v0_5], Discriminator("format_version")
+]
+"""Union of any released model desription"""
+# autogen: stop

bioimageio/spec/notebook/__init__.py

@@ -1,20 +1,20 @@
-# autogen: start
-"""
-implementaions of all released minor versions are available in submodules:
-- notebook v0_2: `bioimageio.spec.notebook.v0_2.NotebookDescr` [user documentation](../../../user_docs/notebook_descr_v0-2.md)
-- notebook v0_3: `bioimageio.spec.notebook.v0_3.NotebookDescr` [user documentation](../../../user_docs/notebook_descr_v0-3.md)
-"""
-from typing import Union
-
-from pydantic import Discriminator
-from typing_extensions import Annotated
-
-from .v0_2 import NotebookDescr as NotebookDescr_v0_2
-from .v0_3 import NotebookDescr as NotebookDescr
-from .v0_3 import NotebookDescr as NotebookDescr_v0_3
-
-AnyNotebookDescr = Annotated[
-    Union[NotebookDescr_v0_2, NotebookDescr_v0_3], Discriminator("format_version")
-]
-"""Union of any released notebook desription"""
-# autogen: stop
+# autogen: start
+"""
+implementaions of all released minor versions are available in submodules:
+- notebook v0_2: `bioimageio.spec.notebook.v0_2.NotebookDescr` [user documentation](../../../user_docs/notebook_descr_v0-2.md)
+- notebook v0_3: `bioimageio.spec.notebook.v0_3.NotebookDescr` [user documentation](../../../user_docs/notebook_descr_v0-3.md)
+"""
+from typing import Union
+
+from pydantic import Discriminator
+from typing_extensions import Annotated
+
+from .v0_2 import NotebookDescr as NotebookDescr_v0_2
+from .v0_3 import NotebookDescr as NotebookDescr
+from .v0_3 import NotebookDescr as NotebookDescr_v0_3
+
+AnyNotebookDescr = Annotated[
+    Union[NotebookDescr_v0_2, NotebookDescr_v0_3], Discriminator("format_version")
+]
+"""Union of any released notebook desription"""
+# autogen: stop

tests/test_package.py

@@ -1,3 +1,4 @@
+import io
 import shutil
 from pathlib import Path
 from tempfile import TemporaryDirectory
@@ -12,6 +13,22 @@ def test_save_bioimageio_package(unet2d_path: Path):
     assert package_path.exists()
 
 
+def test_save_bioimageio_package_to_stream(unet2d_path: Path):
+    from bioimageio.spec import save_bioimageio_package_to_stream
+
+    output_stream = io.BytesIO()
+    stream = save_bioimageio_package_to_stream(unet2d_path, output_stream=output_stream)
+    assert isinstance(stream, io.BytesIO)
+    assert stream == output_stream
+
+
+def test_save_bioimageio_package_to_stream_default(unet2d_path: Path):
+    from bioimageio.spec import save_bioimageio_package_to_stream
+
+    stream = save_bioimageio_package_to_stream(unet2d_path)
+    assert isinstance(stream, io.BytesIO)
+
+
 def test_save_bioimageio_package_as_folder(unet2d_path: Path):
     from bioimageio.spec import load_description, save_bioimageio_package_as_folder