Coverage for /usr/local/lib/python3.12/site-packages/prefect/_internal/pydantic/v2_schema.py: 44%

62 statements  

« prev     ^ index     » next       coverage.py v7.10.6, created at 2025-12-05 10:48 +0000

1import inspect 1a

2import typing 1a

3import typing as t 1a

4 

5import pydantic 1a

6from pydantic import BaseModel as V2BaseModel 1a

7from pydantic import ConfigDict, PydanticUndefinedAnnotation, create_model 1a

8from pydantic.fields import FieldInfo 1a

9from pydantic.type_adapter import TypeAdapter 1a

10 

11from prefect._internal.pydantic.schemas import GenerateEmptySchemaForUserClasses 1a

12 

13 

14def is_v2_model(v: t.Any) -> bool: 1a

15 if isinstance(v, V2BaseModel): 

16 return True 

17 try: 

18 if inspect.isclass(v) and issubclass(v, V2BaseModel): 

19 return True 

20 except TypeError: 

21 pass 

22 

23 return False 

24 

25 

26def is_v2_type(v: t.Any) -> bool: 1a

27 if is_v2_model(v): 

28 return True 

29 

30 try: 

31 return v.__module__.startswith("pydantic.types") 

32 except AttributeError: 

33 return False 

34 

35 

36def has_v2_type_as_param(signature: inspect.Signature) -> bool: 1a

37 parameters = signature.parameters.values() 

38 for p in parameters: 

39 # check if this parameter is a v2 model 

40 if is_v2_type(p.annotation): 

41 return True 

42 

43 # check if this parameter is a collection of types 

44 for v in typing.get_args(p.annotation): 

45 if is_v2_type(v): 

46 return True 

47 return False 

48 

49 

50def process_v2_params( 1a

51 param: inspect.Parameter, 

52 *, 

53 position: int, 

54 docstrings: dict[str, str], 

55 aliases: dict[str, str], 

56) -> tuple[str, t.Any, t.Any]: 

57 """ 

58 Generate a sanitized name, type, and pydantic.Field for a given parameter. 

59 

60 This implementation is exactly the same as the v1 implementation except 

61 that it uses pydantic v2 constructs. 

62 """ 

63 # Pydantic model creation will fail if names collide with the BaseModel type 

64 if hasattr(pydantic.BaseModel, param.name): 64 ↛ 65line 64 didn't jump to line 65 because the condition on line 64 was never true1a

65 name = param.name + "__" 

66 aliases[name] = param.name 

67 else: 

68 name = param.name 1a

69 

70 type_ = t.Any if param.annotation is inspect.Parameter.empty else param.annotation 1a

71 

72 existing_field = param.default if isinstance(param.default, FieldInfo) else None 1a

73 default_value = existing_field.default if existing_field else param.default 1a

74 if existing_field and existing_field.description: 74 ↛ 75line 74 didn't jump to line 75 because the condition on line 74 was never true1a

75 description = existing_field.description 

76 else: 

77 description = docstrings.get(param.name) 1a

78 

79 extra: dict[str, typing.Any] = {} 1a

80 if existing_field and isinstance(existing_field.json_schema_extra, dict): 80 ↛ 82line 80 didn't jump to line 82 because the condition on line 80 was never true1a

81 # this will allow us to merge with the existing `json_schema_extra` 

82 extra.update(existing_field.json_schema_extra) 

83 

84 # still ensure 'position' is always set 

85 extra.setdefault("position", position) 1a

86 

87 field = pydantic.Field( 1a

88 default=... if default_value is param.empty else default_value, 

89 title=param.name, 

90 description=description, 

91 alias=aliases.get(name), 

92 json_schema_extra=extra, 

93 ) 

94 

95 return name, type_, field 1a

96 

97 

98def create_v2_schema( 1a

99 name_: str, 

100 model_cfg: t.Optional[ConfigDict] = None, 

101 model_base: t.Optional[type[V2BaseModel]] = None, 

102 model_fields: t.Optional[dict[str, t.Any]] = None, 

103) -> dict[str, t.Any]: 

104 """ 

105 Create a pydantic v2 model and craft a v1 compatible schema from it. 

106 """ 

107 model_fields = model_fields or {} 1a

108 model = create_model( 1a

109 name_, __config__=model_cfg, __base__=model_base, **model_fields 

110 ) 

111 try: 1a

112 adapter = TypeAdapter(model) 1a

113 except PydanticUndefinedAnnotation as exc: 

114 # in v1 this raises a TypeError, which is handled by parameter_schema 

115 raise TypeError(exc.message) 

116 

117 # root model references under #definitions 

118 schema = adapter.json_schema( 1a

119 by_alias=True, 

120 ref_template="#/definitions/{model}", 

121 schema_generator=GenerateEmptySchemaForUserClasses, 

122 ) 

123 # ensure backwards compatibility by copying $defs into definitions 

124 if "$defs" in schema: 124 ↛ 125line 124 didn't jump to line 125 because the condition on line 124 was never true1a

125 schema["definitions"] = schema["$defs"] 

126 

127 return schema 1a