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, info # pyright: ignore[reportArgumentType] # taken from pydantic docs 

50 ) 

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

52 return ret 

53 

54 

55def _package( 

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

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

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

59 # convert to standard python obj 

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

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

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

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

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

65 if isinstance(source, Path): 

66 unpackaged = source 

67 elif isinstance(source, HttpUrl): 

68 unpackaged = source 

69 elif isinstance(source, RelativeFilePath): 

70 unpackaged = Path(source.path) 

71 elif isinstance(source, AnyUrl): 

72 unpackaged = str(source) 

73 else: 

74 assert_never(source) 

75 

76 if info.mode_is_json(): 

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

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

79 unpackaged = str(unpackaged) 

80 elif isinstance(unpackaged, str): 

81 pass 

82 else: 

83 assert_never(unpackaged) 

84 else: 

85 warnings.warn( 

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

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

88 + "standard python objects" 

89 ) 

90 

91 return unpackaged # return unpackaged file source 

92 

93 fname = extract_file_name(source) 

94 if fname == packaging_context.bioimageio_yaml_file_name: 

95 raise ValueError( 

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

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

98 ) 

99 

100 fsrcs = packaging_context.file_sources 

101 assert not any( 

102 fname.endswith(special) for special in ALL_BIOIMAGEIO_YAML_NAMES 

103 ), fname 

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

105 for i in range(2, 20): 

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

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

108 if ( 

109 alternative_file_name not in fsrcs 

110 or fsrcs[alternative_file_name].source == source 

111 ): 

112 fname = alternative_file_name 

113 break 

114 else: 

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

116 

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

118 return fname 

119 

120 

121include_in_package = PlainSerializer(_package_serializer, when_used="unless-none") 

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

123 

124include_when_packaging = WrapSerializer( 

125 package_file_descr_serializer, when_used="unless-none" 

126) 

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

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

129 

130FileSource_ = Annotated[ 

131 FileSource, 

132 AfterValidator(wo_special_file_name), 

133 include_in_package, 

134] 

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

136 

137FileDescr_ = Annotated[ 

138 FileDescr, AfterValidator(wo_special_file_name), include_when_packaging 

139] 

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