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( 

45 value, 

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

47 ) 

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

49 return ret 

50 

51 

52def _package( 

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

54) -> str | Path | FileName: 

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

56 # convert to standard python obj 

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

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

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

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

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

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

63 unpackaged = source 

64 elif isinstance(source, RelativeFilePath): 

65 unpackaged = Path(source.path) 

66 elif isinstance(source, AnyUrl): 

67 unpackaged = str(source) 

68 else: 

69 assert_never(source) 

70 

71 if info.mode_is_json(): 

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

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

74 unpackaged = str(unpackaged) 

75 elif isinstance(unpackaged, str): 

76 pass 

77 else: 

78 assert_never(unpackaged) 

79 else: 

80 warnings.warn( 

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

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

83 + "standard python objects" 

84 ) 

85 

86 return unpackaged # return unpackaged file source 

87 

88 fname = extract_file_name(source) 

89 if not fname: 

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

91 

92 if fname == packaging_context.bioimageio_yaml_file_name: 

93 raise ValueError( 

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

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

96 ) 

97 

98 if packaging_context.local_files_only: 

99 # skip regular or "relative" URLs 

100 absolute_source = source.absolute() 

101 if isinstance(absolute_source, HttpUrl): 

102 return absolute_source 

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_package = 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_package = Annotated[ 

142 FileDescr, AfterValidator(wo_special_file_name), include_when_packaging 

143] 

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