[Python] When type checking, how to make a pydantic model attribute accept multiple input...

I'm using the Pydantic library, version 2.

I have created a custom class that subclasses int and implements a __get_pydantic_core_schema__ method.

When using this custom class in a Pydantic model, I want the type checker (Basedpyright) to understand that the attribute of the instantiated model will be of the type of my custom class, which works, but also that I can instantiate the model with a str or an int.

However, I don't want to define overloads or override the __init__ method on every model that uses my custom class, as it defeats the purpose and makes the code very verbose, especially if I need to redefine all attributes in every __init__ overload/override.

Here's a minimal reproducible example:

from typing import Any, reveal_type
from pydantic_core import CoreSchema, core_schema
from pydantic import GetCoreSchemaHandler, BaseModel

class MyCustomInt(int):
@classmethod
def __get_pydantic_core_schema__(
cls,
_source_type: Any,
_handler: GetCoreSchemaHandler,
) -> CoreSchema:
def validate_from_int(value: int) -> "MyCustomInt":
if value < 0:
raise ValueError("Cannot be negative")
return cls(value)

from_int_schema = core_schema.chain_schema(
[
core_schema.int_schema(),
core_schema.no_info_plain_validator_function(validate_from_int),
]
)

return core_schema.json_or_python_schema(
json_schema=from_int_schema,
python_schema=core_schema.union_schema(
[
core_schema.is_instance_schema(cls),
from_int_schema,
]
),
serialization=core_schema.plain_serializer_function_ser_schema(
str,
return_schema=core_schema.str_schema(),
when_used='json'
),
)

class MyModel(BaseModel):
value: MyCustomInt

model = MyModel(value="123") # Error: Argument of type "str" cannot be assigned to parameter "value" of type "MyCustomInt" in function "__init__"
reveal_type(model.value) # Revealed type is 'MyCustomInt' - Good!
model = MyModel(value=123) # Error: Argument of type "int" cannot be assigned to parameter "value" of type "MyCustomInt" in function "__init__"
reveal_type(model.value) # Revealed type is 'MyCustomInt' - Good!


To be clear, this works perfectly fine at runtime, and this is purely a type-checking issue.

Is there a way to inform the type checker about the acceptable validation types for MyCustomInt without having to overload or override __init__ in every model?

Continue reading...