diff --git a/mypy/stubgen.py b/mypy/stubgen.py index ece22ba235bf..4d8a1e5f084c 100755 --- a/mypy/stubgen.py +++ b/mypy/stubgen.py @@ -538,6 +538,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: @@ -860,6 +861,8 @@ 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 + if self._inherits_from_pydantic_basemodel(o): + self.processing_pydantic_model = True super().visit_class_def(o) self.dedent() self._vars.pop() @@ -877,6 +880,21 @@ 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 _inherits_from_pydantic_basemodel(self, class_def: ClassDef) -> bool: + """Check if a class directly or indirectly inherits from pydantic.BaseModel""" + for base_type_expr in class_def.base_type_exprs: + if ( + isinstance(base_type_expr, (NameExpr, MemberExpr)) + and self.get_fullname(base_type_expr) == "pydantic.BaseModel" + ): + return True + if self.analyzed and class_def.info: + for base_class in class_def.info.mro: + if base_class.fullname == "pydantic.BaseModel": + return True + return False def get_base_types(self, cdef: ClassDef) -> list[str]: """Get list of base classes for a class.""" @@ -1341,6 +1359,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 "" diff --git a/test-data/unit/stubgen.test b/test-data/unit/stubgen.test index 161f14e8aea7..32020a11bac2 100644 --- a/test-data/unit/stubgen.test +++ b/test-data/unit/stubgen.test @@ -4745,7 +4745,6 @@ class DCMeta(type): ... class DC(metaclass=DCMeta): x: str - [case testIncompleteReturn] from _typeshed import Incomplete @@ -4756,3 +4755,92 @@ def polar(*args, **kwargs) -> Incomplete: from _typeshed import Incomplete def polar(*args, **kwargs) -> Incomplete: ... + +[case testPydanticBaseModel] +import pydantic +from typing import Dict, List, Optional, Union + +class User(pydantic.BaseModel): + id: int + name: str + active: bool = True + optional_field: str | None = None + tags: List[str] = [] + properties: Dict[str, Union[str, int, float, bool]] = {} +[out] +import pydantic + +class User(pydantic.BaseModel): + id: int + name: str + active: bool = ... + optional_field: str | None = ... + tags: list[str] = ... + properties: dict[str, str | int | float | 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 testPydanticBaseModelInheritance_semanal] +from pydantic import BaseModel + +class BaseUser(BaseModel): + id: int + active: bool = True + +class User(BaseUser): + name: str + email: str = '@' +[out] +from pydantic import BaseModel + +class BaseUser(BaseModel): + id: int + active: bool = ... + +class User(BaseUser): + name: str + email: str = ... + +[case testPydanticModelWithMethods] +from pydantic import BaseModel + +class User(BaseModel): + id: int + name: str + + def get_display_name(self) -> str: + return 'a' + + @property + def display_id(self) -> str: + return 'b' +[out] +from pydantic import BaseModel + +class User(BaseModel): + id: int + name: str + def get_display_name(self) -> str: ... + @property + def display_id(self) -> str: ...