Coverage for src / bioimageio / spec / _package.py: 80%
95 statements
« prev ^ index » next coverage.py v7.13.4, created at 2026-02-23 10:51 +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"{source} is invalid: {descr.validation_summary}")
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 exported.validation_summary.display()
273 msg = f"Exported package at '{output_path}' is invalid."
274 if allow_invalid:
275 logger.error(msg)
276 else:
277 raise ValueError(msg)
279 return output_path
282def save_bioimageio_package_to_stream(
283 source: Union[BioimageioYamlSource, ResourceDescr],
284 /,
285 *,
286 compression: int = ZIP_DEFLATED,
287 compression_level: int = 1,
288 output_stream: Union[IO[bytes], None] = None,
289 weights_priority_order: Optional[ # model only
290 Sequence[
291 Literal[
292 "keras_hdf5",
293 "onnx",
294 "pytorch_state_dict",
295 "tensorflow_js",
296 "tensorflow_saved_model_bundle",
297 "torchscript",
298 ]
299 ]
300 ] = None,
301) -> IO[bytes]:
302 """Package a bioimageio resource into a stream.
304 Args:
305 source: bioimageio resource description
306 compression: The numeric constant of compression method.
307 compression_level: Compression level to use when writing files to the archive.
308 See https://docs.python.org/3/library/zipfile.html#zipfile.ZipFile
309 output_stream: stream to write package to
310 weights_priority_order: If given only the first weights format present in the model is included.
311 If none of the prioritized weights formats is found all are included.
313 Note: this function bypasses safety checks and does not load/validate the model after writing.
315 Returns:
316 stream of zipped bioimageio package
317 """
318 if output_stream is None:
319 output_stream = BytesIO()
321 package_content = _prepare_resource_package(
322 source,
323 weights_priority_order=weights_priority_order,
324 )
326 write_zip(
327 output_stream,
328 package_content,
329 compression=compression,
330 compression_level=compression_level,
331 )
333 return output_stream