add dataclass_transform (#1054)

Co-authored-by: Erik De Bonte <erikd@microsoft.com>

4 files changed, 163 insertions, 0 deletions

diff --git a/typing_extensions/CHANGELOG b/typing_extensions/CHANGELOG
index 026ba59..fd1abb9 100644
--- a/typing_extensions/CHANGELOG
+++ b/typing_extensions/CHANGELOG
@@ -1,5 +1,6 @@
 # Release 4.x.x
 
+- Runtime support for PEP 681 and `typing_extensions.dataclass_transform`.
 - `Annotated` can now wrap `ClassVar` and `Final`. Backport from
   bpo-46491. Patch by Gregory Beauregard (@GBeauregard).
 - Add missed `Required` and `NotRequired` to `__all__`. Patch by
diff --git a/typing_extensions/README.rst b/typing_extensions/README.rst
index d5d4128..28081c3 100644
--- a/typing_extensions/README.rst
+++ b/typing_extensions/README.rst
@@ -37,6 +37,7 @@ This module currently contains the following:
 
 - Experimental features
 
+  - ``@dataclass_transform()`` (see PEP 681)
   - ``NotRequired`` (see PEP 655)
   - ``Required`` (see PEP 655)
   - ``Self`` (see PEP 673)
diff --git a/typing_extensions/src/test_typing_extensions.py b/typing_extensions/src/test_typing_extensions.py
index f92bb13..cb32dd0 100644
--- a/typing_extensions/src/test_typing_extensions.py
+++ b/typing_extensions/src/test_typing_extensions.py
@@ -22,6 +22,7 @@ from typing_extensions import NoReturn, ClassVar, Final, IntVar, Literal, Type,
 from typing_extensions import TypeAlias, ParamSpec, Concatenate, ParamSpecArgs, ParamSpecKwargs, TypeGuard
 from typing_extensions import Awaitable, AsyncIterator, AsyncContextManager, Required, NotRequired
 from typing_extensions import Protocol, runtime, runtime_checkable, Annotated, overload, final, is_typeddict
+from typing_extensions import dataclass_transform
 try:
     from typing_extensions import get_type_hints
 except ImportError:
@@ -2345,6 +2346,83 @@ class FinalDecoratorTests(BaseTestCase):
         self.assertIs(True, Methods.cached.__final__)
 
 
+class DataclassTransformTests(BaseTestCase):
+    def test_decorator(self):
+        def create_model(*, frozen: bool = False, kw_only: bool = True):
+            return lambda cls: cls
+
+        decorated = dataclass_transform(kw_only_default=True, order_default=False)(create_model)
+
+        class CustomerModel:
+            id: int
+
+        self.assertIs(decorated, create_model)
+        self.assertEqual(
+            decorated.__dataclass_transform__,
+            {
+                "eq_default": True,
+                "order_default": False,
+                "kw_only_default": True,
+                "field_descriptors": (),
+            }
+        )
+        self.assertIs(
+            decorated(frozen=True, kw_only=False)(CustomerModel),
+            CustomerModel
+        )
+
+    def test_base_class(self):
+        class ModelBase:
+            def __init_subclass__(cls, *, frozen: bool = False): ...
+
+        Decorated = dataclass_transform(eq_default=True, order_default=True)(ModelBase)
+
+        class CustomerModel(Decorated, frozen=True):
+            id: int
+
+        self.assertIs(Decorated, ModelBase)
+        self.assertEqual(
+            Decorated.__dataclass_transform__,
+            {
+                "eq_default": True,
+                "order_default": True,
+                "kw_only_default": False,
+                "field_descriptors": (),
+            }
+        )
+        self.assertIsSubclass(CustomerModel, Decorated)
+
+    def test_metaclass(self):
+        class Field: ...
+
+        class ModelMeta(type):
+            def __new__(
+                cls, name, bases, namespace, *, init: bool = True,
+            ):
+                return super().__new__(cls, name, bases, namespace)
+
+        Decorated = dataclass_transform(
+            order_default=True, field_descriptors=(Field,)
+        )(ModelMeta)
+
+        class ModelBase(metaclass=Decorated): ...
+
+        class CustomerModel(ModelBase, init=False):
+            id: int
+
+        self.assertIs(Decorated, ModelMeta)
+        self.assertEqual(
+            Decorated.__dataclass_transform__,
+            {
+                "eq_default": True,
+                "order_default": True,
+                "kw_only_default": False,
+                "field_descriptors": (Field,),
+            }
+        )
+        self.assertIsInstance(CustomerModel, Decorated)
+
+
 class AllTests(BaseTestCase):
 
     def test_typing_extensions_includes_standard(self):
diff --git a/typing_extensions/src/typing_extensions.py b/typing_extensions/src/typing_extensions.py
index dab53d4..bcf1695 100644
--- a/typing_extensions/src/typing_extensions.py
+++ b/typing_extensions/src/typing_extensions.py
@@ -70,6 +70,7 @@ __all__ = [
 
     # One-off things.
     'Annotated',
+    'dataclass_transform',
     'final',
     'IntVar',
     'is_typeddict',
@@ -2341,3 +2342,85 @@ else:
 
     Required = _Required(_root=True)
     NotRequired = _NotRequired(_root=True)
+
+if hasattr(typing, 'dataclass_transform'):
+    dataclass_transform = typing.dataclass_transform
+else:
+    def dataclass_transform(
+        *,
+        eq_default: bool = True,
+        order_default: bool = False,
+        kw_only_default: bool = False,
+        field_descriptors: typing.Tuple[
+            typing.Union[typing.Type[typing.Any], typing.Callable[..., typing.Any]],
+            ...
+        ] = (),
+    ) -> typing.Callable[[T], T]:
+        """Decorator that marks a function, class, or metaclass as providing
+        dataclass-like behavior.
+
+        Example:
+
+            from typing_extensions import dataclass_transform
+
+            _T = TypeVar("_T")
+
+            # Used on a decorator function
+            @dataclass_transform()
+            def create_model(cls: type[_T]) -> type[_T]:
+                ...
+                return cls
+
+            @create_model
+            class CustomerModel:
+                id: int
+                name: str
+
+            # Used on a base class
+            @dataclass_transform()
+            class ModelBase: ...
+
+            class CustomerModel(ModelBase):
+                id: int
+                name: str
+
+            # Used on a metaclass
+            @dataclass_transform()
+            class ModelMeta(type): ...
+
+            class ModelBase(metaclass=ModelMeta): ...
+
+            class CustomerModel(ModelBase):
+                id: int
+                name: str
+
+        Each of the ``CustomerModel`` classes defined in this example will now
+        behave similarly to a dataclass created with the ``@dataclasses.dataclass``
+        decorator. For example, the type checker will synthesize an ``__init__``
+        method.
+
+        The arguments to this decorator can be used to customize this behavior:
+        - ``eq_default`` indicates whether the ``eq`` parameter is assumed to be
+          True or False if it is omitted by the caller.
+        - ``order_default`` indicates whether the ``order`` parameter is
+          assumed to be True or False if it is omitted by the caller.
+        - ``kw_only_default`` indicates whether the ``kw_only`` parameter is
+          assumed to be True or False if it is omitted by the caller.
+        - ``field_descriptors`` specifies a static list of supported classes
+          or functions, that describe fields, similar to ``dataclasses.field()``.
+
+        At runtime, this decorator records its arguments in the
+        ``__dataclass_transform__`` attribute on the decorated object.
+
+        See PEP 681 for details.
+
+        """
+        def decorator(cls_or_fn):
+            cls_or_fn.__dataclass_transform__ = {
+                "eq_default": eq_default,
+                "order_default": order_default,
+                "kw_only_default": kw_only_default,
+                "field_descriptors": field_descriptors,
+            }
+            return cls_or_fn
+        return decorator