Coverage for bioimageio/spec/_internal/io_packaging.py: 79%

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 PlainSerializer, 

13 SerializationInfo, 

14 SerializerFunctionWrapHandler, 

15 WrapSerializer, 

16) 

17from typing_extensions import ( 

18 Annotated, 

19 assert_never, 

20) 

21 

22from .io import ( 

23 FileDescr, 

24 FileSource, 

25 RelativeFilePath, 

26 extract_file_name, 

27 wo_special_file_name, 

28) 

29from .io_basics import ( 

30 ALL_BIOIMAGEIO_YAML_NAMES, 

31 FileName, 

32 Sha256, 

33) 

34from .packaging_context import packaging_context_var 

35from .url import HttpUrl 

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 fname == packaging_context.bioimageio_yaml_file_name: 

96 raise ValueError( 

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

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

99 ) 

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 = PlainSerializer(_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_ = 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_ = Annotated[ 

139 FileDescr, AfterValidator(wo_special_file_name), include_when_packaging 

140] 

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