Override dataclass_transform behavior for RootModel (#8163)

pydantic/mypy.py

@@ -860,8 +860,10 @@ def add_initializer(
                 use_alias=use_alias,
                 is_settings=is_settings,
             )
-            if is_root_model:
+
+            if is_root_model and MYPY_VERSION_TUPLE <= (1, 0, 1):
                 # convert root argument to positional argument
+                # This is needed because mypy support for `dataclass_transform` isn't complete on 1.0.1
                 args[0].kind = ARG_POS if args[0].kind == ARG_NAMED else ARG_OPT
 
             if is_settings:

pydantic/root_model.py

@@ -8,24 +8,34 @@
 from pydantic_core import PydanticUndefined
 
 from . import PydanticUserError
-from ._internal import _repr
+from ._internal import _model_construction, _repr
 from .main import BaseModel, _object_setattr
 
 if typing.TYPE_CHECKING:
     from typing import Any
 
-    from typing_extensions import Literal
+    from typing_extensions import Literal, dataclass_transform
 
-    Model = typing.TypeVar('Model', bound='BaseModel')
+    from .fields import Field as PydanticModelField
+
+    # dataclass_transform could be applied to RootModel directly, but `ModelMetaclass`'s dataclass_transform
+    # takes priority (at least with pyright). We trick type checkers into thinking we apply dataclass_transform
+    # on a new metaclass.
+    @dataclass_transform(kw_only_default=False, field_specifiers=(PydanticModelField,))
+    class _RootModelMetaclass(_model_construction.ModelMetaclass):
+        ...
 
+    Model = typing.TypeVar('Model', bound='BaseModel')
+else:
+    _RootModelMetaclass = _model_construction.ModelMetaclass
 
 __all__ = ('RootModel',)
 
 
 RootModelRootType = typing.TypeVar('RootModelRootType')
 
 
-class RootModel(BaseModel, typing.Generic[RootModelRootType]):
+class RootModel(BaseModel, typing.Generic[RootModelRootType], metaclass=_RootModelMetaclass):
     """Usage docs: https://docs.pydantic.dev/2.6/concepts/models/#rootmodel-and-custom-root-types
 
     A Pydantic `BaseModel` for the root object of the model.

tests/mypy/outputs/1.1.1/mypy-default_ini/root_models.py

@@ -0,0 +1,24 @@
+from typing import List
+
+from pydantic import RootModel
+
+
+class Pets1(RootModel[List[str]]):
+    pass
+
+
+Pets2 = RootModel[List[str]]
+
+
+class Pets3(RootModel):
+# MYPY: error: Missing type parameters for generic type "RootModel"  [type-arg]
+    root: List[str]
+
+
+pets1 = Pets1(['dog', 'cat'])
+pets2 = Pets2(['dog', 'cat'])
+pets3 = Pets3(['dog', 'cat'])
+
+
+class Pets4(RootModel[List[str]]):
+    pets: List[str]

tests/mypy/outputs/1.1.1/mypy-plugin_ini/root_models.py

@@ -0,0 +1,25 @@
+from typing import List
+
+from pydantic import RootModel
+
+
+class Pets1(RootModel[List[str]]):
+    pass
+
+
+Pets2 = RootModel[List[str]]
+
+
+class Pets3(RootModel):
+# MYPY: error: Missing type parameters for generic type "RootModel"  [type-arg]
+    root: List[str]
+
+
+pets1 = Pets1(['dog', 'cat'])
+pets2 = Pets2(['dog', 'cat'])
+pets3 = Pets3(['dog', 'cat'])
+
+
+class Pets4(RootModel[List[str]]):
+    pets: List[str]
+# MYPY: error: Only `root` is allowed as a field of a `RootModel`  [pydantic-field]

tests/mypy/outputs/1.1.1/pyproject-default_toml/root_models.py

@@ -0,0 +1,24 @@
+from typing import List
+
+from pydantic import RootModel
+
+
+class Pets1(RootModel[List[str]]):
+    pass
+
+
+Pets2 = RootModel[List[str]]
+
+
+class Pets3(RootModel):
+# MYPY: error: Missing type parameters for generic type "RootModel"  [type-arg]
+    root: List[str]
+
+
+pets1 = Pets1(['dog', 'cat'])
+pets2 = Pets2(['dog', 'cat'])
+pets3 = Pets3(['dog', 'cat'])
+
+
+class Pets4(RootModel[List[str]]):
+    pets: List[str]

tests/mypy/outputs/1.4.1/mypy-default_ini/root_models.py

@@ -0,0 +1,24 @@
+from typing import List
+
+from pydantic import RootModel
+
+
+class Pets1(RootModel[List[str]]):
+    pass
+
+
+Pets2 = RootModel[List[str]]
+
+
+class Pets3(RootModel):
+# MYPY: error: Missing type parameters for generic type "RootModel"  [type-arg]
+    root: List[str]
+
+
+pets1 = Pets1(['dog', 'cat'])
+pets2 = Pets2(['dog', 'cat'])
+pets3 = Pets3(['dog', 'cat'])
+
+
+class Pets4(RootModel[List[str]]):
+    pets: List[str]

tests/mypy/outputs/1.4.1/pyproject-default_toml/root_models.py

@@ -0,0 +1,24 @@
+from typing import List
+
+from pydantic import RootModel
+
+
+class Pets1(RootModel[List[str]]):
+    pass
+
+
+Pets2 = RootModel[List[str]]
+
+
+class Pets3(RootModel):
+# MYPY: error: Missing type parameters for generic type "RootModel"  [type-arg]
+    root: List[str]
+
+
+pets1 = Pets1(['dog', 'cat'])
+pets2 = Pets2(['dog', 'cat'])
+pets3 = Pets3(['dog', 'cat'])
+
+
+class Pets4(RootModel[List[str]]):
+    pets: List[str]

tests/mypy/test_mypy.py

@@ -69,6 +69,13 @@ def build(self) -> List[Union[Tuple[str, str], Any]]:
         'success.py',
         pytest.mark.skipif(MYPY_VERSION_TUPLE > (1, 0, 1), reason='Need to handle some more things for mypy >=1.1.1'),
     ).build()
+    + MypyCasesBuilder(
+        ['mypy-default.ini', 'pyproject-default.toml'],
+        'root_models.py',
+        pytest.mark.skipif(
+            MYPY_VERSION_TUPLE < (1, 1, 1), reason='`dataclass_transform` only supported on mypy >= 1.1.1'
+        ),
+    ).build()
     + MypyCasesBuilder(
         'mypy-default.ini', ['plugin_success.py', 'plugin_success_baseConfig.py', 'metaclass_args.py']
     ).build()