8000 Add obj to methods · pydantic/pydantic@c3e2b75 · GitHub
[go: up one dir, main page]
Skip to content

Commit c3e2b75

Browse files
adriangbadriangb
committed
Add obj to methods
1 parent 5bd636b commit c3e2b75

File tree

4 files changed

+61
-40
lines changed

4 files changed

+61
-40
lines changed

‎pydantic/_internal/_generate_schema.py‎

Lines changed: 58 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -326,46 +326,61 @@ def _current_generate_schema(self) -> GenerateSchema:
326326
def _arbitrary_types(self) -> bool:
327327
return self._config_wrapper.arbitrary_types_allowed
328328

329-
def literal_schema(self, values: list[Any]) -> CoreSchema:
329+
def literal_schema(self, tp: Any, values: list[Any]) -> CoreSchema:
330330
return core_schema.literal_schema(values)
331331

332-
def any_schema(self) -> CoreSchema:
332+
def any_schema(self, tp: Any) -> CoreSchema:
333333
return core_schema.any_schema()
334334

335-
def str_schema(self) -> CoreSchema:
335+
def str_schema(self, tp: Any) -> CoreSchema:
336336
return core_schema.str_schema()
337337

338-
def bytes_schema(self) -> CoreSchema:
338+
def bytes_schema(
339+
self,
340+
tp: Any,
341+
) -> CoreSchema:
339342
return core_schema.bytes_schema()
340343

341-
def int_schema(self) -> CoreSchema:
344+
def int_schema(
345+
self,
346+
tp: Any,
347+
) -> CoreSchema:
342348
return core_schema.int_schema()
343349

344-
def float_schema(self) -> CoreSchema:
350+
def float_schema(
351+
self,
352+
tp: Any,
353+
) -> CoreSchema:
345354
return core_schema.float_schema()
346355

347-
def bool_schema(self) -> CoreSchema:
356+
def bool_schema(
357+
self,
358+
tp: Any,
359+
) -> CoreSchema:
348360
return core_schema.bool_schema()
349361

350-
def none_schema(self) -> CoreSchema:
362+
def none_schema(
363+
self,
364+
tp: Any,
365+
) -> CoreSchema:
351366
return core_schema.none_schema()
352367

353-
def list_schema(self, items_type: Any) -> CoreSchema:
368+
def list_schema(self, tp: Any, items_type: Any) -> CoreSchema:
354369
return core_schema.list_schema(self.generate_schema(items_type))
355370

356-
def dict_schema(self, keys_type: Any, values_type: Any) -> CoreSchema:
371+
def dict_schema(self, tp: Any, keys_type: Any, values_type: Any) -> CoreSchema:
357372
return core_schema.dict_schema(self.generate_schema(keys_type), self.generate_schema(values_type))
358373

359-
def set_schema(self, items_type: Any) -> CoreSchema:
374+
def set_schema(self, tp: Any, items_type: Any) -> CoreSchema:
360375
return core_schema.set_schema(self.generate_schema(items_type))
361376

362-
def frozenset_schema(self, items_type: Any) -> CoreSchema:
377+
def frozenset_schema(self, tp: Any, items_type: Any) -> CoreSchema:
363378
return core_schema.frozenset_schema(self.generate_schema(items_type))
364379

365-
def tuple_variable_schema(self, items_type: Any) -> CoreSchema:
380+
def tuple_variable_schema(self, tp: Any, items_type: Any) -> CoreSchema:
366381
return core_schema.tuple_variable_schema(self.generate_schema(items_type))
367382

368-
def tuple_positional_schema(self, items_types: list[Any]) -> CoreSchema:
383+
def tuple_positional_schema(self, tp: Any, items_types: list[Any]) -> CoreSchema:
369384
items_schemas = [self.generate_schema(items_type) for items_type in items_types]
370385
return core_schema.tuple_positional_schema(items_schemas)
371386

@@ -667,6 +682,15 @@ def _get_first_arg_or_any(self, obj: Any) -> Any:
667682
return Any
668683
return args[0]
669684

685+
def _get_first_two_args_or_any(self, obj: Any) -> tuple[Any, Any]:
686+
args = self._get_args_resolving_forward_refs(obj)
687+
if not args:
688+
return (Any, Any)
689+
if len(args) < 2:
690+
origin = get_origin(obj)
691+
raise TypeError(f'Expected two type arguments for {origin}, got 1')
692+
return args[0], args[1]
693+
670694
def _generate_schema(self, obj: Any) -> core_schema.CoreSchema:
671695
"""Recursively generate a pydantic-core schema for any supported python type."""
672696
if isinstance(obj, dict):
@@ -703,29 +727,29 @@ def match_type(self, obj: Any) -> core_schema.CoreSchema: # noqa: C901
703727
as they get requested and we figure out what the right API for them is.
704728
"""
705729
if obj is str:
706-
return self.str_schema()
730+
return self.str_schema(obj)
707731
elif obj is bytes:
708-
return self.bytes_schema()
732+
return self.bytes_schema(obj)
709733
elif obj is int:
710-
return self.int_schema()
734+
return self.int_schema(obj)
711735
elif obj is float:
712-
return self.float_schema()
736+
return self.float_schema(obj)
713737
elif obj is bool:
714-
return self.bool_schema()
738+
return self.bool_schema(obj)
715739
elif obj is Any or obj is object:
716-
return self.any_schema()
740+
return self.any_schema(obj)
717741
elif obj is None or obj is _typing_extra.NoneType:
718-
return self.none_schema()
742+
return self.none_schema(obj)
719743
elif obj in TUPLE_TYPES:
720744
return self._tuple_schema(obj)
721745
elif obj in LIST_TYPES:
722-
return self.list_schema(self._get_first_arg_or_any(obj))
746+
return self.list_schema(obj, self._get_first_arg_or_any(obj))
723747
elif obj in SET_TYPES:
724-
return self.set_schema(self._get_first_arg_or_any(obj))
748+
return self.set_schema(obj, self._get_first_arg_or_any(obj))
725749
elif obj in FROZEN_SET_TYPES:
726-
return self.frozenset_schema(self._get_first_arg_or_any(obj))
750+
return self.frozenset_schema(obj, self._get_first_arg_or_any(obj))
727751
elif obj in DICT_TYPES:
728-
return self.dict_schema(*(self._get_args_resolving_forward_refs(obj) or (Any, Any)))
752+
return self.dict_schema(obj, *self._get_first_two_args_or_any(obj))
729753
elif isinstance(obj, TypeAliasType):
730754
return self._type_alias_type_schema(obj)
731755
elif obj == type:
@@ -751,10 +775,7 @@ def match_type(self, obj: Any) -> core_schema.CoreSchema: # noqa: C901
751775
if< A2BF /span> obj is Final:
752776
return core_schema.any_schema()
753777
return self.generate_schema(
754-
self._get_args_resolving_forward_refs(
755-
obj,
756-
required=True,
757-
)[0]
778+
self._get_first_arg_or_any(obj),
758779
)
759780
elif isinstance(obj, (FunctionType, LambdaType, MethodType, partial)):
760781
return self._callable_schema(obj)
@@ -801,13 +822,13 @@ def _match_generic_type(self, obj: Any, origin: Any) -> CoreSchema: # noqa: C90
801822
elif origin in TUPLE_TYPES:
802823
return self._tuple_schema(obj)
803824
elif origin in LIST_TYPES:
804-
return self.list_schema(self._get_first_arg_or_any(obj))
825+
return self.list_schema(obj, self._get_first_arg_or_any(obj))
805826
elif origin in SET_TYPES:
806-
return self.set_schema(self._get_first_arg_or_any(obj))
827+
return self.set_schema(obj, self._get_first_arg_or_any(obj))
807828
elif origin in FROZEN_SET_TYPES:
808-
return self.frozenset_schema(self._get_first_arg_or_any(obj))
829+
return self.frozenset_schema(obj, self._get_first_arg_or_any(obj))
809830
elif origin in DICT_TYPES:
810-
return self.dict_schema(*(self._get_args_resolving_forward_refs(obj) or (Any, Any)))
831+
return self.dict_schema(obj, *self._get_first_two_args_or_any(obj))
811832
elif is_typeddict(origin):
812833
return self._typed_dict_schema(obj, origin)
813834
elif origin in (typing.Type, type):
@@ -1005,7 +1026,7 @@ def _literal_schema(self, literal_type: Any) -> CoreSchema:
10051026
"""Generate schema for a Literal."""
10061027
expected = _typing_extra.all_literal_values(literal_type)
10071028
assert expected, f'literal "expected" cannot be empty, obj={literal_type}'
1008-
return self.literal_schema(expected)
1029+
return self.literal_schema(literal_type, expected)
10091030

10101031
def _typed_dict_schema(self, typed_dict_cls: Any, origin: Any) -> core_schema.CoreSchema:
10111032
"""Generate schema for a TypedDict.
@@ -1168,16 +1189,16 @@ def _tuple_schema(self, tuple_type: Any) -> core_schema.CoreSchema:
11681189
return core_schema.tuple_positional_schema([])
11691190
elif params[-1] is Ellipsis:
11701191
if len(params) == 2:
1171-
return self.tuple_variable_schema(params[0])
1192+
return self.tuple_variable_schema(tuple_type, params[0])
11721193
else:
11731194
# TODO: something like https://github.com/pydantic/pydantic/issues/5952
11741195
raise ValueError('Variable tuples can only have one type')
11751196
elif len(params) == 1 and params[0] == ():
11761197
# special case for `Tuple[()]` which means `Tuple[]` - an empty tuple
11771198
# NOTE: This conditional can be removed when we drop support for Python 3.10.
1178-
return self.tuple_positional_schema([])
1199+
return self.tuple_positional_schema(tuple_type, [])
11791200
else:
1180-
return self.tuple_positional_schema(list(params))
1201+
return self.tuple_positional_schema(tuple_type, list(params))
11811202

11821203
def _type_schema(self) -> core_schema.CoreSchema:
11831204
return core_schema.custom_error_schema(

‎tests/test_dataclasses.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2457,7 +2457,7 @@ class B:
24572457

24582458
def test_schema_generator() -> None:
24592459
class LaxStrGenerator(GenerateSchema):
2460-
def str_schema(self) -> CoreSchema:
2460+
def str_schema(self, _obj: Any) -> CoreSchema:
24612461
return core_schema.no_info_plain_validator_function(str)
24622462

24632463
@pydantic.dataclasses.dataclass

‎tests/test_main.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2799,7 +2799,7 @@ class Model3(BaseModel):
27992799

28002800
def test_schema_generator_customize_type() -> None:
28012801
class LaxStrGenerator(GenerateSchema):
2802-
def str_schema(self) -> CoreSchema:
2802+
def str_schema(self, obj: Any) -> CoreSchema:
28032803
return core_schema.no_info_plain_validator_function(str)
28042804

28052805
class Model(BaseModel):

‎tests/test_types_typeddict.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -868,7 +868,7 @@ class Model(TypedDict):
868868

869869
def test_schema_generator() -> None:
870870
class LaxStrGenerator(GenerateSchema):
871-
def str_schema(self) -> CoreSchema:
871+
def str_schema(self, obj: Any) -> CoreSchema:
872872
return core_schema.no_info_plain_validator_function(str)
873873

874874
class Model(TypedDict):

0 commit comments

Comments
 (0)
0