From 8efaf8a9b1317174dc2712991129e162c27bf059 Mon Sep 17 00:00:00 2001 From: fynnbe Date: Fri, 17 Nov 2023 10:25:21 +0100 Subject: [PATCH] update io functions --- bioimageio/spec/__init__.py | 2 +- bioimageio/spec/_internal/io_utils.py | 4 +- bioimageio/spec/_internal/types/__init__.py | 2 +- bioimageio/spec/_io.py | 6 ++- bioimageio/spec/_package.py | 53 +++++++++++---------- bioimageio/spec/utils.py | 1 + 6 files changed, 37 insertions(+), 31 deletions(-) create mode 100644 bioimageio/spec/utils.py diff --git a/bioimageio/spec/__init__.py b/bioimageio/spec/__init__.py index ab6d35cee..efbeb9e07 100644 --- a/bioimageio/spec/__init__.py +++ b/bioimageio/spec/__init__.py @@ -12,7 +12,7 @@ from bioimageio.spec._description import validate_format as validate_format from bioimageio.spec._internal.constants import VERSION from bioimageio.spec._io import load_description as load_description -from bioimageio.spec._io import save_description as save_description +from bioimageio.spec._io import save_bioimageio_yaml_only as save_bioimageio_yaml_only from bioimageio.spec.application import AnyApplication as AnyApplication from bioimageio.spec.application import Application as Application from bioimageio.spec.collection import AnyCollection as AnyCollection diff --git a/bioimageio/spec/_internal/io_utils.py b/bioimageio/spec/_internal/io_utils.py index 391d36449..8ef3fc446 100644 --- a/bioimageio/spec/_internal/io_utils.py +++ b/bioimageio/spec/_internal/io_utils.py @@ -10,7 +10,7 @@ from zipfile import ZipFile, is_zipfile import pooch -from pydantic import AnyUrl, DirectoryPath, FilePath, HttpUrl, TypeAdapter +from pydantic import AnyUrl, DirectoryPath, FilePath, HttpUrl, NewPath, TypeAdapter from ruamel.yaml import YAML from typing_extensions import NotRequired, Unpack @@ -87,7 +87,7 @@ def read_yaml(file: Union[FilePath, TextIO]) -> YamlValue: yaml.load(f) -def write_yaml(content: YamlValue, /, file: Union[FilePath, TextIO]): +def write_yaml(content: YamlValue, /, file: Union[NewPath, FilePath, TextIO]): if isinstance(file, Path): cm = file.open("w", encoding="utf-8") else: diff --git a/bioimageio/spec/_internal/types/__init__.py b/bioimageio/spec/_internal/types/__init__.py index b07e14e54..39a91b0fc 100644 --- a/bioimageio/spec/_internal/types/__init__.py +++ b/bioimageio/spec/_internal/types/__init__.py @@ -65,7 +65,6 @@ BeforeValidator(lambda s: s.replace("×", "·").replace("*", "·").replace(" ", "·") if isinstance(s, str) else s), ], ) -Unit = Union[Literal["arbitrary unit", "px"], SiUnit] # types as loaded from YAML 1.2 (with ruamel.yaml) YamlLeafValue = Union[bool, date, datetime, float, int, str, None] @@ -81,3 +80,4 @@ # derived types DatasetId = NewType("DatasetId", ResourceId) BioimageioYamlContent = Dict[str, YamlValue] +BioimageioYamlSource = Union[PermissiveFileSource, BioimageioYamlContent] diff --git a/bioimageio/spec/_io.py b/bioimageio/spec/_io.py index 07a8db95b..c966ee543 100644 --- a/bioimageio/spec/_io.py +++ b/bioimageio/spec/_io.py @@ -1,6 +1,6 @@ from typing import Literal, TextIO, Union, cast -from pydantic import FilePath +from pydantic import FilePath, NewPath from bioimageio.spec import ResourceDescription from bioimageio.spec._description import InvalidDescription, build_description, dump_description @@ -30,7 +30,9 @@ def load_description( ) -def save_description(rd: Union[ResourceDescription, BioimageioYamlContent], /, file: Union[FilePath, TextIO]): +def save_bioimageio_yaml_only( + rd: Union[ResourceDescription, BioimageioYamlContent], /, file: Union[NewPath, FilePath, TextIO] +): if isinstance(rd, ResourceDescriptionBase): content = dump_description(rd) else: diff --git a/bioimageio/spec/_package.py b/bioimageio/spec/_package.py index 5cdcd9f34..1e184447a 100644 --- a/bioimageio/spec/_package.py +++ b/bioimageio/spec/_package.py @@ -6,7 +6,7 @@ from typing import Any, Dict, Literal, Optional, Sequence, Tuple, Union, cast from zipfile import ZIP_DEFLATED -from pydantic import AnyUrl, DirectoryPath, FilePath, HttpUrl +from pydantic import AnyUrl, DirectoryPath, FilePath, HttpUrl, NewPath from bioimageio.spec import model from bioimageio.spec._description import InvalidDescription, ResourceDescription, build_description, dump_description @@ -14,21 +14,24 @@ from bioimageio.spec._internal.constants import IN_PACKAGE_MESSAGE from bioimageio.spec._internal.io_utils import ( BIOIMAGEIO_YAML, - PermissiveFileSource, download, open_bioimageio_yaml, write_yaml, write_zip, ) -from bioimageio.spec._internal.types import BioimageioYamlContent, FileName, RelativeFilePath, YamlValue +from bioimageio.spec._internal.types import ( + BioimageioYamlContent, + BioimageioYamlSource, + FileName, + RelativeFilePath, + YamlValue, +) from bioimageio.spec._internal.types._file_source import extract_file_name from bioimageio.spec._internal.utils import nest_dict_with_narrow_first_key from bioimageio.spec._internal.validation_context import ValidationContext from bioimageio.spec.model.v0_4 import WeightsFormat from bioimageio.spec.summary import Loc -BioimageioYamlSource = Union[PermissiveFileSource, ResourceDescription, BioimageioYamlContent] - def fill_resource_package_content( package_content: Dict[Loc, Union[HttpUrl, RelativeFilePath]], @@ -131,7 +134,7 @@ def get_resource_package_content( def _prepare_resource_package( - rdf_source: BioimageioYamlSource, + source: Union[BioimageioYamlSource, ResourceDescription], /, *, weights_priority_order: Optional[Sequence[WeightsFormat]] = None, @@ -144,18 +147,18 @@ def _prepare_resource_package( 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. """ - if isinstance(rdf_source, ResourceDescriptionBase): - rd = rdf_source + if isinstance(source, ResourceDescriptionBase): + rd = source _ctxt = rd._internal_validation_context # pyright: ignore[reportPrivateUsage] context = ValidationContext(root=_ctxt["root"], file_name=_ctxt["file_name"]) - elif isinstance(rdf_source, dict): + elif isinstance(source, dict): context = ValidationContext() rd = build_description( - rdf_source, + source, context=context, ) else: - rdf = open_bioimageio_yaml(rdf_source) + rdf = open_bioimageio_yaml(source) context = ValidationContext(root=rdf.original_root, file_name=rdf.original_file_name) rd = build_description( rdf.content, @@ -163,7 +166,7 @@ def _prepare_resource_package( ) if isinstance(rd, InvalidDescription): - raise ValueError(f"{rdf_source} is invalid: {rd.validation_summaries[0]}") + raise ValueError(f"{source} is invalid: {rd.validation_summaries[0]}") package_content = get_resource_package_content(rd, weights_priority_order=weights_priority_order) @@ -177,11 +180,11 @@ def _prepare_resource_package( return local_package_content -def write_package_as_folder( - rdf_source: BioimageioYamlSource, +def save_bioimageio_package_as_folder( + source: Union[BioimageioYamlSource, ResourceDescription], /, *, - output_path: Optional[DirectoryPath] = None, + output_path: Union[NewPath, DirectoryPath, None] = None, weights_priority_order: Optional[ # model only Sequence[ Literal[ @@ -198,16 +201,16 @@ def write_package_as_folder( """Write the content of a bioimage.io resource package to a folder. Args: - rd: bioimage.io resource description + source: bioimageio resource description output_path: file path 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. Returns: - path to zipped bioimage.io package in BIOIMAGEIO_CACHE_PATH or 'output_path' + directory path to bioimageio package folder """ package_content = _prepare_resource_package( - rdf_source, + source, weights_priority_order=weights_priority_order, ) if output_path is None: @@ -224,13 +227,13 @@ def write_package_as_folder( return output_path -def write_package( - rdf_source: BioimageioYamlSource, +def save_bioimageio_package( + source: Union[BioimageioYamlSource, ResourceDescription], /, *, compression: int = ZIP_DEFLATED, compression_level: int = 1, - output_path: Optional[Path] = None, + output_path: Union[NewPath, FilePath, None] = None, weights_priority_order: Optional[ # model only Sequence[ Literal[ @@ -244,10 +247,10 @@ def write_package( ] ] = None, ) -> FilePath: - """Package a bioimage.io resource as a zip file. + """Package a bioimageio resource as a zip file. Args: - rd: bioimage.io resource description + 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 @@ -256,10 +259,10 @@ def write_package( If none of the prioritized weights formats is found all are included. Returns: - path to zipped bioimage.io package in BIOIMAGEIO_CACHE_PATH or 'output_path' + path to zipped bioimageio package """ package_content = _prepare_resource_package( - rdf_source, + source, weights_priority_order=weights_priority_order, ) if output_path is None: diff --git a/bioimageio/spec/utils.py b/bioimageio/spec/utils.py new file mode 100644 index 000000000..538379ff9 --- /dev/null +++ b/bioimageio/spec/utils.py @@ -0,0 +1 @@ +from bioimageio.spec._internal.io_utils import download as download