Coverage for src/bioimageio/spec/_internal/io_packaging.py: 74%

1from __future__ import annotations 

2 

3import warnings 

4from pathlib import Path 

5 

6from pydantic import ( 

7 AnyUrl, 

8 SerializationInfo, 

9 SerializerFunctionWrapHandler, 

10 WrapSerializer, 

11) 

12from typing_extensions import ( 

13 Annotated, 

14 assert_never, 

15) 

16 

17from .io import ( 

18 FileDescr, 

19 FileSource, 

20 RelativeFilePath, 

21 extract_file_name, 

22 wo_special_file_name, 

23) 

24from .io_basics import ( 

25 ALL_BIOIMAGEIO_YAML_NAMES, 

26 FileName, 

27 Sha256, 

28) 

29from .packaging_context import packaging_context_var 

30from .url import HttpUrl 

31from .utils import PrettyPlainSerializer 

32from .validator_annotations import AfterValidator 

33 

34 

35def _package_serializer( 

36 source: FileSource, info: SerializationInfo 

37) -> str | Path | FileName: 

38 return _package(source, info, None) 

39 

40 

41def package_file_descr_serializer( 

42 value: FileDescr, handler: SerializerFunctionWrapHandler, info: SerializationInfo 

43): 

44 ret = handler(value) 

45 ret["source"] = _package(value.source, info, sha256=value.sha256) 

46 return ret 

47 

48 

49def _package( 

50 source: FileSource, info: SerializationInfo, sha256: Sha256 | None 

51) -> str | Path | FileName: 

52 if (packaging_context := packaging_context_var.get()) is None: 

53 # convert to standard python obj 

54 # note: pydantic keeps returning Rootmodels (here `HttpUrl`) as-is, but if 

55 # this function returns one RootModel, paths are "further serialized" by 

56 # returning the 'root' attribute, which is incorrect. 

57 # see https://github.com/pydantic/pydantic/issues/8963 

58 # TODO: follow up on https://github.com/pydantic/pydantic/issues/8963 

59 if isinstance(source, (Path, HttpUrl)): 

60 unpackaged = source 

61 elif isinstance(source, RelativeFilePath): 

62 unpackaged = Path(source.path) 

63 elif isinstance(source, AnyUrl): 

64 unpackaged = str(source) 

65 else: 

66 assert_never(source) 

67 

68 if info.mode_is_json(): 

69 # convert to json value # TODO: remove and let pydantic do this? 

70 if isinstance(unpackaged, (Path, HttpUrl)): 

71 unpackaged = str(unpackaged) 

72 elif isinstance(unpackaged, str): 

73 pass 

74 else: 

75 assert_never(unpackaged) 

76 else: 

77 warnings.warn( 

78 "dumping with mode='python' is currently not fully supported for " 

79 + "fields that are included when packaging; returned objects are " 

80 + "standard python objects" 

81 ) 

82 

83 return unpackaged # return unpackaged file source 

84 

85 fname = extract_file_name(source) 

86 if not fname: 

87 raise ValueError(f"Got empty file name for source {source}") 

88 

89 if fname == packaging_context.bioimageio_yaml_file_name: 

90 raise ValueError( 

91 f"Reserved file name '{packaging_context.bioimageio_yaml_file_name}' " 

92 + "not allowed for a file to be packaged" 

93 ) 

94 

95 if packaging_context.local_files_only: 

96 # skip regular or "relative" URLs 

97 absolute_source = source.absolute() 

98 if isinstance(absolute_source, HttpUrl): 

99 return absolute_source 

100 

101 fsrcs = packaging_context.file_sources 

102 assert not any(fname.endswith(special) for special in ALL_BIOIMAGEIO_YAML_NAMES), ( 

103 fname 

104 ) 

105 if fname in fsrcs and fsrcs[fname].source != source: 

106 for i in range(2, 20): 

107 fn, *ext = fname.split(".") 

108 alternative_file_name = ".".join([f"{fn}_{i}", *ext]) 

109 if ( 

110 alternative_file_name not in fsrcs 

111 or fsrcs[alternative_file_name].source == source 

112 ): 

113 fname = alternative_file_name 

114 break 

115 else: 

116 raise ValueError(f"Too many file name clashes for {fname}") 

117 

118 fsrcs[fname] = FileDescr(source=source, sha256=sha256) 

119 return fname 

120 

121 

122include_in_package = PrettyPlainSerializer(_package_serializer, when_used="unless-none") 

123"""DEPRECATED serializer for `source` fields without corresponding `sha256` field.""" 

124 

125include_when_packaging = WrapSerializer( 

126 package_file_descr_serializer, when_used="unless-none" 

127) 

128"""Pydantic serializer that marks the annotated `FileDescr` to be included when packaging 

129(saving a bioimageio zip package).""" 

130 

131FileSource_package = Annotated[ 

132 FileSource, 

133 AfterValidator(wo_special_file_name), 

134 include_in_package, 

135] 

136"""A file source that is included when packaging the resource.""" 

137 

138FileDescr_package = Annotated[ 

139 FileDescr, AfterValidator(wo_special_file_name), include_when_packaging 

140] 

141"""A `FileDescr` whose **source** is included when packaging the resource."""