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