Coverage for src/bioimageio/spec/_package.py: 83%
95 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-18 09:17 +0000
1from __future__ import annotations
3import collections.abc
4import shutil
5from io import BytesIO
6from pathlib import Path
7from tempfile import NamedTemporaryFile, mkdtemp
8from typing import IO, Literal, Sequence
9from zipfile import ZIP_DEFLATED
11from loguru import logger
12from pydantic import DirectoryPath, FilePath, NewPath
14from ._description import InvalidDescr, ResourceDescr, build_description
15from ._internal.common_nodes import ResourceDescrBase
16from ._internal.io import (
17 BioimageioYamlContent,
18 BioimageioYamlSource,
19 FileDescr,
20 RelativeFilePath,
21 ensure_is_valid_bioimageio_yaml_name,
22)
23from ._internal.io_basics import (
24 BIOIMAGEIO_YAML,
25 AbsoluteFilePath,
26 BytesReader,
27 FileName,
28 ZipPath,
29)
30from ._internal.io_utils import open_bioimageio_yaml, write_yaml, write_zip
31from ._internal.packaging_context import PackagingContext
32from ._internal.url import HttpUrl
33from ._internal.utils import get_os_friendly_file_name
34from ._internal.validation_context import get_validation_context
35from ._internal.warning_levels import ERROR
36from ._io import load_description
37from .model.v0_4 import WeightsFormat
40def get_resource_package_content(
41 rd: ResourceDescr,
42 /,
43 *,
44 bioimageio_yaml_file_name: FileName = BIOIMAGEIO_YAML,
45 weights_priority_order: Sequence[WeightsFormat] | None = None, # model only
46) -> dict[FileName, HttpUrl | AbsoluteFilePath | BioimageioYamlContent | ZipPath]:
47 """DEPRECATED in favor of get_package_content: Get the content of a bioimage.io resource package."""
48 ret: dict[
49 FileName, HttpUrl | AbsoluteFilePath | BioimageioYamlContent | ZipPath
50 ] = {}
51 for k, v in get_package_content(
52 rd,
53 bioimageio_yaml_file_name=bioimageio_yaml_file_name,
54 weights_priority_order=weights_priority_order,
55 ).items():
56 if isinstance(v, FileDescr):
57 if isinstance(v.source, (Path, RelativeFilePath)):
58 ret[k] = v.source.absolute()
59 else:
60 ret[k] = v.source
62 else:
63 ret[k] = v
65 return ret
68def get_package_content(
69 rd: ResourceDescr,
70 /,
71 *,
72 bioimageio_yaml_file_name: FileName = BIOIMAGEIO_YAML,
73 weights_priority_order: Sequence[WeightsFormat] | None = None, # model only
74 local_files_only: bool = False,
75) -> dict[FileName, FileDescr | BioimageioYamlContent]:
76 """
77 Args:
78 rd: resource description
79 bioimageio_yaml_file_name: RDF file name
80 weights_priority_order: (for model resources only)
81 If given, only the first weights format present in the model is included.
82 If none of the prioritized weights formats is found a ValueError is raised.
83 local_files_only: If True, only local files are included in the package content. If False, remote files are also included.
84 """
85 os_friendly_name = get_os_friendly_file_name(rd.name)
86 bioimageio_yaml_file_name = bioimageio_yaml_file_name.format(
87 name=os_friendly_name, type=rd.type
88 )
90 bioimageio_yaml_file_name = ensure_is_valid_bioimageio_yaml_name(
91 bioimageio_yaml_file_name
92 )
93 content: dict[FileName, FileDescr] = {}
94 with PackagingContext(
95 bioimageio_yaml_file_name=bioimageio_yaml_file_name,
96 file_sources=content,
97 weights_priority_order=weights_priority_order,
98 local_files_only=local_files_only,
99 ):
100 rdf_content: BioimageioYamlContent = rd.model_dump(
101 mode="json", exclude_unset=True
102 )
104 _ = rdf_content.pop("rdf_source", None)
106 return {**content, bioimageio_yaml_file_name: rdf_content}
109def _prepare_resource_package(
110 source: BioimageioYamlSource | ResourceDescr,
111 /,
112 *,
113 weights_priority_order: Sequence[WeightsFormat] | None = None,
114 local_files_only: bool = False,
115) -> dict[FileName, BioimageioYamlContent | BytesReader]:
116 """Prepare to package a resource description; downloads all required files.
118 Args:
119 source: A bioimage.io resource description (as file, raw YAML content or description class)
120 context: validation context
121 weights_priority_order: If given only the first weights format present in the model is included.
122 If none of the prioritized weights formats is found all are included.
123 local_files_only: If True, only local files are included in the package. If False, remote files are also included.
124 """
125 context = get_validation_context()
126 bioimageio_yaml_file_name = context.file_name
127 if isinstance(source, ResourceDescrBase):
128 descr = source
129 elif isinstance(source, collections.abc.Mapping):
130 descr = build_description(source)
131 else:
132 opened = open_bioimageio_yaml(source)
133 bioimageio_yaml_file_name = opened.original_file_name
134 context = context.replace(
135 root=opened.original_root, file_name=opened.original_file_name
136 )
137 with context:
138 descr = build_description(opened.content)
140 if isinstance(descr, InvalidDescr):
141 raise ValueError(f"Resource description is invalid:\n{descr.get_reason()}")
143 with context:
144 package_content = get_package_content(
145 descr,
146 bioimageio_yaml_file_name=bioimageio_yaml_file_name or BIOIMAGEIO_YAML,
147 weights_priority_order=weights_priority_order,
148 local_files_only=local_files_only,
149 )
151 return {
152 k: v if isinstance(v, collections.abc.Mapping) else v.get_reader()
153 for k, v in package_content.items()
154 }
157def save_bioimageio_package_as_folder(
158 source: BioimageioYamlSource | ResourceDescr,
159 /,
160 *,
161 output_path: NewPath | DirectoryPath | None = None,
162 weights_priority_order: Sequence[
163 Literal[
164 "keras_hdf5",
165 "onnx",
166 "pytorch_state_dict",
167 "tensorflow_js",
168 "tensorflow_saved_model_bundle",
169 "torchscript",
170 ]
171 ]
172 | None = None,
173 local_files_only: bool = False,
174) -> DirectoryPath:
175 """Write the content of a bioimage.io resource package to a folder.
177 Args:
178 source: bioimageio resource description
179 output_path: file path to write package to
180 weights_priority_order: If given only the first weights format present in the model is included.
181 If none of the prioritized weights formats is found all are included.
182 local_files_only: If True, only local files are included in the package. If False, remote files are also included.
184 Returns:
185 directory path to bioimageio package folder
186 """
187 package_content = _prepare_resource_package(
188 source,
189 weights_priority_order=weights_priority_order,
190 local_files_only=local_files_only,
191 )
192 if output_path is None:
193 output_path = Path(mkdtemp())
194 else:
195 output_path = Path(output_path)
197 output_path.mkdir(exist_ok=True, parents=True)
198 for name, src in package_content.items():
199 if not name:
200 raise ValueError("got empty file name in package content")
202 if isinstance(src, collections.abc.Mapping):
203 write_yaml(src, output_path / name)
204 elif (
205 isinstance(src.original_root, Path)
206 and src.original_root / src.original_file_name
207 == (output_path / name).resolve()
208 ):
209 logger.debug(
210 f"Not copying {src.original_root / src.original_file_name} to itself."
211 )
212 else:
213 if isinstance(src.original_root, Path):
214 logger.debug(
215 f"Copying from path {src.original_root / src.original_file_name} to {output_path / name}."
216 )
217 else:
218 logger.debug(
219 f"Copying {src.original_root}/{src.original_file_name} to {output_path / name}."
220 )
221 with (output_path / name).open("wb") as dest:
222 _ = shutil.copyfileobj(src, dest)
224 return output_path
227def save_bioimageio_package(
228 source: BioimageioYamlSource | ResourceDescr,
229 /,
230 *,
231 compression: int = ZIP_DEFLATED,
232 compression_level: int = 1,
233 output_path: NewPath | FilePath | None = None,
234 weights_priority_order: Sequence[
235 Literal[
236 "keras_hdf5",
237 "onnx",
238 "pytorch_state_dict",
239 "tensorflow_js",
240 "tensorflow_saved_model_bundle",
241 "torchscript",
242 ]
243 ]
244 | None = None,
245 allow_invalid: bool = False,
246 local_files_only: bool = False,
247) -> FilePath:
248 """Package a bioimageio resource as a zip file.
250 Args:
251 source: bioimageio resource description
252 compression: The numeric constant of compression method.
253 compression_level: Compression level to use when writing files to the archive.
254 See https://docs.python.org/3/library/zipfile.html#zipfile.ZipFile
255 output_path: file path to write package to
256 weights_priority_order: If given only the first weights format present in the model is included.
257 If none of the prioritized weights formats is found all are included.
258 allow_invalid: If True, do not raise an error if the exported package is invalid, but log an error instead.
259 local_files_only: If True, only local files are included in the package. If False, remote files are also included.
261 Returns:
262 path to zipped bioimageio package
263 """
264 package_content = _prepare_resource_package(
265 source,
266 weights_priority_order=weights_priority_order,
267 local_files_only=local_files_only,
268 )
269 if output_path is None:
270 output_path = Path(
271 NamedTemporaryFile(suffix=".bioimageio.zip", delete=False).name # ruff: ignore[SIM115]
272 )
273 else:
274 output_path = Path(output_path)
276 write_zip(
277 output_path,
278 package_content,
279 compression=compression,
280 compression_level=compression_level,
281 )
282 with get_validation_context().replace(warning_level=ERROR):
283 if isinstance((exported := load_description(output_path)), InvalidDescr):
284 msg = f"Exported package at '{output_path}' is invalid:\n{exported.get_reason()}"
285 if allow_invalid:
286 logger.error(msg)
287 else:
288 raise ValueError(msg)
290 return output_path
293def save_bioimageio_package_to_stream(
294 source: BioimageioYamlSource | ResourceDescr,
295 /,
296 *,
297 compression: int = ZIP_DEFLATED,
298 compression_level: int = 1,
299 output_stream: IO[bytes] | None = None,
300 weights_priority_order: Sequence[
301 Literal[
302 "keras_hdf5",
303 "onnx",
304 "pytorch_state_dict",
305 "tensorflow_js",
306 "tensorflow_saved_model_bundle",
307 "torchscript",
308 ]
309 ]
310 | None = None,
311 local_files_only: bool = False,
312) -> IO[bytes]:
313 """Package a bioimageio resource into a stream.
315 Args:
316 source: bioimageio resource description
317 compression: The numeric constant of compression method.
318 compression_level: Compression level to use when writing files to the archive.
319 See https://docs.python.org/3/library/zipfile.html#zipfile.ZipFile
320 output_stream: stream to write package to
321 weights_priority_order: If given only the first weights format present in the model is included.
322 If none of the prioritized weights formats is found all are included.
323 local_files_only: If True, only local files are included in the package. If False, remote files are also included.
325 Note: this function bypasses safety checks and does not load/validate the model after writing.
327 Returns:
328 stream of zipped bioimageio package
329 """
330 if output_stream is None:
331 output_stream = BytesIO()
333 package_content = _prepare_resource_package(
334 source,
335 weights_priority_order=weights_priority_order,
336 local_files_only=local_files_only,
337 )
339 write_zip(
340 output_stream,
341 package_content,
342 compression=compression,
343 compression_level=compression_level,
344 )
346 return output_stream