Add support for Pydantic models in stubgen

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 ""

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: ...