Coverage for src/bioimageio/spec/_package.py: 83%
94 statements
« prev ^ index » next coverage.py v7.14.1, created at 2026-06-15 15:08 +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
38def get_resource_package_content(
39 rd: ResourceDescr,
40 /,
41 *,
42 bioimageio_yaml_file_name: FileName = BIOIMAGEIO_YAML,
43 weights_priority_order: Optional[Sequence[WeightsFormat]] = None, # model only
44) -> Dict[FileName, Union[HttpUrl, AbsoluteFilePath, BioimageioYamlContent, ZipPath]]:
45 """DEPRECATED in favor of get_package_content: Get the content of a bioimage.io resource package."""
46 ret: Dict[
47 FileName, Union[HttpUrl, AbsoluteFilePath, BioimageioYamlContent, ZipPath]
48 ] = {}
49 for k, v in get_package_content(
50 rd,
51 bioimageio_yaml_file_name=bioimageio_yaml_file_name,
52 weights_priority_order=weights_priority_order,
53 ).items():
54 if isinstance(v, FileDescr):
55 if isinstance(v.source, (Path, RelativeFilePath)):
56 ret[k] = v.source.absolute()
57 else:
58 ret[k] = v.source
60 else:
61 ret[k] = v
63 return ret
66def get_package_content(
67 rd: ResourceDescr,
68 /,
69 *,
70 bioimageio_yaml_file_name: FileName = BIOIMAGEIO_YAML,
71 weights_priority_order: Optional[Sequence[WeightsFormat]] = None, # model only
72 local_files_only: bool = False,
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 local_files_only: If True, only local files are included in the package content. If False, remote files are also included.
82 """
83 os_friendly_name = get_os_friendly_file_name(rd.name)
84 bioimageio_yaml_file_name = bioimageio_yaml_file_name.format(
85 name=os_friendly_name, type=rd.type
86 )
88 bioimageio_yaml_file_name = ensure_is_valid_bioimageio_yaml_name(
89 bioimageio_yaml_file_name
90 )
91 content: Dict[FileName, FileDescr] = {}
92 with PackagingContext(
93 bioimageio_yaml_file_name=bioimageio_yaml_file_name,
94 file_sources=content,
95 weights_priority_order=weights_priority_order,
96 local_files_only=local_files_only,
97 ):
98 rdf_content: BioimageioYamlContent = rd.model_dump(
99 mode="json", exclude_unset=True
100 )
102 _ = rdf_content.pop("rdf_source", None)
104 return {**content, bioimageio_yaml_file_name: rdf_content}
107def _prepare_resource_package(
108 source: Union[BioimageioYamlSource, ResourceDescr],
109 /,
110 *,
111 weights_priority_order: Optional[Sequence[WeightsFormat]] = None,
112 local_files_only: bool = False,
113) -> Dict[FileName, Union[BioimageioYamlContent, BytesReader]]:
114 """Prepare to package a resource description; downloads all required files.
116 Args:
117 source: A bioimage.io resource description (as file, raw YAML content or description class)
118 context: validation context
119 weights_priority_order: If given only the first weights format present in the model is included.
120 If none of the prioritized weights formats is found all are included.
121 local_files_only: If True, only local files are included in the package. If False, remote files are also included.
122 """
123 context = get_validation_context()
124 bioimageio_yaml_file_name = context.file_name
125 if isinstance(source, ResourceDescrBase):
126 descr = source
127 elif isinstance(source, collections.abc.Mapping):
128 descr = build_description(source)
129 else:
130 opened = open_bioimageio_yaml(source)
131 bioimageio_yaml_file_name = opened.original_file_name
132 context = context.replace(
133 root=opened.original_root, file_name=opened.original_file_name
134 )
135 with context:
136 descr = build_description(opened.content)
138 if isinstance(descr, InvalidDescr):
139 raise ValueError(f"Resource description is invalid:\n{descr.get_reason()}")
141 with context:
142 package_content = get_package_content(
143 descr,
144 bioimageio_yaml_file_name=bioimageio_yaml_file_name or BIOIMAGEIO_YAML,
145 weights_priority_order=weights_priority_order,
146 local_files_only=local_files_only,
147 )
149 return {
150 k: v if isinstance(v, collections.abc.Mapping) else v.get_reader()
151 for k, v in package_content.items()
152 }
155def save_bioimageio_package_as_folder(
156 source: Union[BioimageioYamlSource, ResourceDescr],
157 /,
158 *,
159 output_path: Union[NewPath, DirectoryPath, None] = None,
160 weights_priority_order: Optional[ # model only
161 Sequence[
162 Literal[
163 "keras_hdf5",
164 "onnx",
165 "pytorch_state_dict",
166 "tensorflow_js",
167 "tensorflow_saved_model_bundle",
168 "torchscript",
169 ]
170 ]
171 ] = None,
172 local_files_only: bool = False,
173) -> DirectoryPath:
174 """Write the content of a bioimage.io resource package to a folder.
176 Args:
177 source: bioimageio resource description
178 output_path: file path to write package to
179 weights_priority_order: If given only the first weights format present in the model is included.
180 If none of the prioritized weights formats is found all are included.
181 local_files_only: If True, only local files are included in the package. If False, remote files are also included.
183 Returns:
184 directory path to bioimageio package folder
185 """
186 package_content = _prepare_resource_package(
187 source,
188 weights_priority_order=weights_priority_order,
189 local_files_only=local_files_only,
190 )
191 if output_path is None:
192 output_path = Path(mkdtemp())
193 else:
194 output_path = Path(output_path)
196 output_path.mkdir(exist_ok=True, parents=True)
197 for name, src in package_content.items():
198 if not name:
199 raise ValueError("got empty file name in package content")
201 if isinstance(src, collections.abc.Mapping):
202 write_yaml(src, output_path / name)
203 elif (
204 isinstance(src.original_root, Path)
205 and src.original_root / src.original_file_name
206 == (output_path / name).resolve()
207 ):
208 logger.debug(
209 f"Not copying {src.original_root / src.original_file_name} to itself."
210 )
211 else:
212 if isinstance(src.original_root, Path):
213 logger.debug(
214 f"Copying from path {src.original_root / src.original_file_name} to {output_path / name}."
215 )
216 else:
217 logger.debug(
218 f"Copying {src.original_root}/{src.original_file_name} to {output_path / name}."
219 )
220 with (output_path / name).open("wb") as dest:
221 _ = shutil.copyfileobj(src, dest)
223 return output_path
226def save_bioimageio_package(
227 source: Union[BioimageioYamlSource, ResourceDescr],
228 /,
229 *,
230 compression: int = ZIP_DEFLATED,
231 compression_level: int = 1,
232 output_path: Union[NewPath, FilePath, None] = None,
233 weights_priority_order: Optional[ # model only
234 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,
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
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: Union[BioimageioYamlSource, ResourceDescr],
295 /,
296 *,
297 compression: int = ZIP_DEFLATED,
298 compression_level: int = 1,
299 output_stream: Union[IO[bytes], None] = None,
300 weights_priority_order: Optional[ # model only
301 Sequence[
302 Literal[
303 "keras_hdf5",
304 "onnx",
305 "pytorch_state_dict",
306 "tensorflow_js",
307 "tensorflow_saved_model_bundle",
308 "torchscript",
309 ]
310 ]
311 ] = None,
312 local_files_only: bool = False,
313) -> IO[bytes]:
314 """Package a bioimageio resource into a stream.
316 Args:
317 source: bioimageio resource description
318 compression: The numeric constant of compression method.
319 compression_level: Compression level to use when writing files to the archive.
320 See https://docs.python.org/3/library/zipfile.html#zipfile.ZipFile
321 output_stream: stream to write package to
322 weights_priority_order: If given only the first weights format present in the model is included.
323 If none of the prioritized weights formats is found all are included.
324 local_files_only: If True, only local files are included in the package. If False, remote files are also included.
326 Note: this function bypasses safety checks and does not load/validate the model after writing.
328 Returns:
329 stream of zipped bioimageio package
330 """
331 if output_stream is None:
332 output_stream = BytesIO()
334 package_content = _prepare_resource_package(
335 source,
336 weights_priority_order=weights_priority_order,
337 local_files_only=local_files_only,
338 )
340 write_zip(
341 output_stream,
342 package_content,
343 compression=compression,
344 compression_level=compression_level,
345 )
347 return output_stream