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

1from __future__ import annotations 

2 

3import warnings 

4from pathlib import Path 

5from typing import ( 

6 Optional, 

7 Union, 

8) 

9 

10from pydantic import ( 

11 AnyUrl, 

12 SerializationInfo, 

13 SerializerFunctionWrapHandler, 

14 WrapSerializer, 

15) 

16from typing_extensions import ( 

17 Annotated, 

18 assert_never, 

19) 

20 

21from .io import ( 

22 FileDescr, 

23 FileSource, 

24 RelativeFilePath, 

25 extract_file_name, 

26 wo_special_file_name, 

27) 

28from .io_basics import ( 

29 ALL_BIOIMAGEIO_YAML_NAMES, 

30 FileName, 

31 Sha256, 

32) 

33from .packaging_context import packaging_context_var 

34from .url import HttpUrl 

35from .utils import PrettyPlainSerializer 

36from .validator_annotations import AfterValidator 

37 

38 

39def _package_serializer( 

40 source: FileSource, info: SerializationInfo 

41) -> Union[str, Path, FileName]: 

42 return _package(source, info, None) 

43 

44 

45def package_file_descr_serializer( 

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

47): 

48 ret = handler( 

49 value, 

50 info, # pyright: ignore[reportArgumentType] # taken from pydantic docs 

51 ) 

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

53 return ret 

54 

55 

56def _package( 

57 source: FileSource, info: SerializationInfo, sha256: Optional[Sha256] 

58) -> Union[str, Path, FileName]: 

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

60 # convert to standard python obj 

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

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

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

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

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

66 if isinstance(source, Path): 

67 unpackaged = source 

68 elif isinstance(source, HttpUrl): 

69 unpackaged = source 

70 elif isinstance(source, RelativeFilePath): 

71 unpackaged = Path(source.path) 

72 elif isinstance(source, AnyUrl): 

73 unpackaged = str(source) 

74 else: 

75 assert_never(source) 

76 

77 if info.mode_is_json(): 

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

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

80 unpackaged = str(unpackaged) 

81 elif isinstance(unpackaged, str): 

82 pass 

83 else: 

84 assert_never(unpackaged) 

85 else: 

86 warnings.warn( 

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

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

89 + "standard python objects" 

90 ) 

91 

92 return unpackaged # return unpackaged file source 

93 

94 fname = extract_file_name(source) 

95 if not fname: 

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

97 

98 if fname == packaging_context.bioimageio_yaml_file_name: 

99 raise ValueError( 

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

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

102 ) 

103 

104 fsrcs = packaging_context.file_sources 

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

106 fname 

107 ) 

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

109 for i in range(2, 20): 

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

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

112 if ( 

113 alternative_file_name not in fsrcs 

114 or fsrcs[alternative_file_name].source == source 

115 ): 

116 fname = alternative_file_name 

117 break 

118 else: 

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

120 

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

122 return fname 

123 

124 

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

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

127 

128include_when_packaging = WrapSerializer( 

129 package_file_descr_serializer, when_used="unless-none" 

130) 

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

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

133 

134FileSource_ = Annotated[ 

135 FileSource, 

136 AfterValidator(wo_special_file_name), 

137 include_in_package, 

138] 

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

140 

141FileDescr_ = Annotated[ 

142 FileDescr, AfterValidator(wo_special_file_name), include_when_packaging 

143] 

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