Skip to content

Add support for Pydantic models in stubgen #19095

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 13 commits into
base: master
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from 7 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 13 additions & 0 deletions mypy/stubgen.py
Original file line number Diff line number Diff line change
Expand Up @@ -491,6 +491,7 @@ def __init__(
self.processing_enum = False
self.processing_dataclass = False
self.dataclass_field_specifier: tuple[str, ...] = ()
self.processing_pydantic_model = False

@property
def _current_class(self) -> ClassDef | None:
Expand Down Expand Up @@ -808,6 +809,14 @@ def visit_class_def(self, o: ClassDef) -> None:
if self.analyzed and (spec := find_dataclass_transform_spec(o)):
self.processing_dataclass = True
self.dataclass_field_specifier = spec.field_specifiers
is_pydantic_model = False
for base_type_expr in o.base_type_exprs:
if isinstance(base_type_expr, (NameExpr, MemberExpr)) and self.get_fullname(
base_type_expr
).endswith("BaseModel"):
is_pydantic_model = True
break
self.processing_pydantic_model = is_pydantic_model
super().visit_class_def(o)
self.dedent()
self._vars.pop()
Expand All @@ -825,6 +834,7 @@ def visit_class_def(self, o: ClassDef) -> None:
self.dataclass_field_specifier = ()
self._class_stack.pop(-1)
self.processing_enum = False
self.processing_pydantic_model = False

def get_base_types(self, cdef: ClassDef) -> list[str]:
"""Get list of base classes for a class."""
Expand Down Expand Up @@ -1289,6 +1299,9 @@ def get_assign_initializer(self, rvalue: Expression) -> str:
if not (isinstance(rvalue, TempNode) and rvalue.no_rhs):
return " = ..."
# TODO: support other possible cases, where initializer is important
if self.processing_pydantic_model:
if not (isinstance(rvalue, TempNode) and rvalue.no_rhs):
return " = ..."

# By default, no initializer is required:
return ""
Expand Down
118 changes: 118 additions & 0 deletions test-data/unit/stubgen.test
Original file line number Diff line number Diff line change
Expand Up @@ -4718,3 +4718,121 @@ class DCMeta(type): ...

class DC(metaclass=DCMeta):
x: str

[case testPydanticBaseModel]
import pydantic

class User(pydantic.BaseModel):
id: int
name: str
active: bool = True
optional_field: str | None = None
[out]
import pydantic

class User(pydantic.BaseModel):
id: int
name: str
active: bool = ...
optional_field: str | None = ...

[case testPydanticBaseModelWithAnnotationsOnly]
import pydantic

class ConfigSettings(pydantic.BaseModel):
# Fields without initialization
db_name: str
port: int
debug: bool
[out]
import pydantic

class ConfigSettings(pydantic.BaseModel):
db_name: str
port: int
debug: bool

[case testPydanticNestedBaseModel]
from pydantic import BaseModel

class Address(BaseModel):
street: str
city: str

class User(BaseModel):
name: str
age: int
address: Address | None = None
[out]
from pydantic import BaseModel

class Address(BaseModel):
street: str
city: str

class User(BaseModel):
name: str
age: int
address: Address | None = ...

[case testPydanticBaseModelComplex]
from pydantic import BaseModel
from typing import Dict, List, Optional, Union

class Item(BaseModel):
name: str
description: Optional[str] = None
tags: List[str] = []
properties: Dict[str, Union[str, int, float, bool]] = {}
[out]
from pydantic import BaseModel

class Item(BaseModel):
name: str
description: str | None = ...
tags: list[str] = ...
properties: dict[str, str | int | float | bool] = ...

[case testPydanticBaseModelInheritance]
from pydantic import BaseModel

class BaseUser(BaseModel):
id: int
active: bool = True

class User(BaseUser):
name: str
email: str
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
email: str
email: str = '@'

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There is an issue connected with this change. Currently, the code only sets processing_pydantic_model = True when the direct base class is pydantic.BaseModel, but it doesn't handle cases where a class inherits from another class that inherits from pydantic.BaseModel. I tried a few ideas but they aren't working. Could you please give me a clue how to tackle the problem?

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hi @sobolevn, do you have any idea how to detect the indirect inheritance from pydantic.BaseModel? Besides that, do you know why ellipses aren't propagated for the attributes with default values in general? I guess it might be relevant for all classes, not only pydantic descendants.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What is "indirect inheritance"?

Copy link
Author

@teplandr teplandr Jul 3, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe it's not the correct term, but I mean that if we have:

class BaseUser(BaseModel):
    pass

class User(BaseUser):
    pass

BaseUser directly inherited from BaseModel (and base_type_expr contains this info), whereas User indirectly inherited from BaseModel (and base_type_expr doesn't contain info about BaseModel).

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We can inspect the mro of a TypeInfo

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hi @sobolevn! I pushed the changes with Pydantic inheritance detection via mro. Unfortunately, mro contains the following list: [<TypeInfo main.User>, <TypeInfo main.BaseUser>, <TypeInfo builtins.object>] and no pydantic.BaseModel. Could you help me to identify the reason?

[out]
from pydantic import BaseModel

class BaseUser(BaseModel):
id: int
active: bool = ...

class User(BaseUser):
name: str
email: str

[case testPydanticModelWithMethods]
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We can just add some methods to existing models. Let's reduce the amount of tests without reducing test features.

from pydantic import BaseModel

class User(BaseModel):
id: int
name: str

def get_display_name(self) -> str:
return f"User {self.name}"

@property
def display_id(self) -> str:
return f"ID: {self.id}"
[out]
from pydantic import BaseModel

class User(BaseModel):
id: int
name: str
def get_display_name(self) -> str: ...
@property
def display_id(self) -> str: ...
pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy