1
1
from typing import Callable , Optional , Tuple , Union
2
2
3
3
from mypy import nodes
4
- from mypy .plugin import ClassDefContext , FunctionContext , Plugin , Type
4
+ from mypy .plugin import ClassDefContext , Plugin
5
5
from pydantic .mypy import PydanticModelTransformer , PydanticPlugin
6
6
7
7
MODEL_METACLASS_FULLNAME = 'pydantic_xml.model.XmlModelMeta'
@@ -21,38 +21,6 @@ def get_metaclass_hook(self, fullname: str) -> Optional[Callable[[ClassDefContex
21
21
return self ._pydantic_model_metaclass_marker_callback
22
22
return super ().get_metaclass_hook (fullname )
23
23
24
- def get_function_hook (self , fullname : str ) -> Optional [Callable [[FunctionContext ], Type ]]:
25
- sym = self .lookup_fully_qualified (fullname )
26
- if sym and sym .fullname == ATTR_FULLNAME :
27
- return self ._attribute_callback
28
- elif sym and sym .fullname == ELEMENT_FULLNAME :
29
- return self ._element_callback
30
- elif sym and sym .fullname == WRAPPED_FULLNAME :
31
- return self ._wrapped_callback
32
-
33
- return super ().get_function_hook (fullname )
34
-
35
- def _attribute_callback (self , ctx : FunctionContext ) -> Type :
36
- return super ()._pydantic_field_callback (self ._pop_first_args (ctx , 2 ))
37
-
38
- def _element_callback (self , ctx : FunctionContext ) -> Type :
39
- return super ()._pydantic_field_callback (self ._pop_first_args (ctx , 4 ))
40
-
41
- def _wrapped_callback (self , ctx : FunctionContext ) -> Type :
42
- return super ()._pydantic_field_callback (self ._pop_first_args (ctx , 4 ))
43
-
44
- def _pop_first_args (self , ctx : FunctionContext , num : int ) -> FunctionContext :
45
- return FunctionContext (
46
- arg_types = ctx .arg_types [num :],
47
- arg_kinds = ctx .arg_kinds [num :],
48
- callee_arg_names = ctx .callee_arg_names [num :],
49
- arg_names = ctx .arg_names [num :],
50
- default_return_type = ctx .default_return_type ,
51
- args = ctx .args [num :],
52
- context = ctx .context ,
53
- api = ctx .api ,
54
- )
55
-
56
24
def _pydantic_model_class_maker_callback (self , ctx : ClassDefContext ) -> bool :
57
25
transformer = PydanticXmlModelTransformer (ctx .cls , ctx .reason , ctx .api , self .plugin_config )
58
26
return transformer .transform ()
@@ -100,3 +68,43 @@ def get_alias_info(stmt: nodes.AssignmentStmt) -> Tuple[Union[str, None], bool]:
100
68
return None , True
101
69
102
70
return PydanticModelTransformer .get_alias_info (stmt )
71
+
72
+ @staticmethod
73
+ def get_strict (stmt : nodes .AssignmentStmt ) -> Optional [bool ]:
74
+ expr = stmt .rvalue
75
+ if (
76
+ isinstance (expr , nodes .CallExpr ) and
77
+ isinstance (expr .callee , nodes .RefExpr ) and
78
+ expr .callee .fullname in ENTITIES_FULLNAME
79
+ ):
80
+ for arg , name in zip (expr .args , expr .arg_names ):
81
+ if name != 'strict' :
82
+ continue
83
+ if isinstance (arg , nodes .NameExpr ):
84
+ if arg .fullname == 'builtins.True' :
85
+ return True
86
+ elif arg .fullname == 'builtins.False' :
87
+ return False
88
+ return None
89
+
90
+ return PydanticModelTransformer .get_strict (stmt )
91
+
92
+ @staticmethod
93
+ def is_field_frozen (stmt : nodes .AssignmentStmt ) -> bool :
94
+ expr = stmt .rvalue
95
+ if isinstance (expr , nodes .TempNode ):
96
+ return False
97
+
98
+ if not (
99
+ isinstance (expr , nodes .CallExpr ) and
100
+ isinstance (expr .callee , nodes .RefExpr ) and
101
+ expr .callee .fullname in ENTITIES_FULLNAME
102
+ ):
103
+ return False
104
+
105
+ for i , arg_name in enumerate (expr .arg_names ):
106
+ if arg_name == 'frozen' :
107
+ arg = expr .args [i ]
108
+ return isinstance (arg , nodes .NameExpr ) and arg .fullname == 'builtins.True'
109
+
110
+ return PydanticModelTransformer .is_field_frozen (stmt )
0 commit comments