consolidate functions for finding dataclass transform specs

python · JukkaL · Feb 8, 2023 · Jan 30, 2023 · Feb 1, 2023 · Feb 3, 2023
commit 12c8c74b81389f48cd5f6cff429614ff159e30ec
diff --git a/mypy/plugins/dataclasses.py b/mypy/plugins/dataclasses.py
@@ -14,13 +14,11 @@
     ARG_STAR,
     ARG_STAR2,
     MDEF,
-    SYMBOL_FUNCBASE_TYPES,
     Argument,
     AssignmentStmt,
     CallExpr,
     Context,
     DataclassTransformSpec,
-    Decorator,
     Expression,
     JsonDict,
     NameExpr,
@@ -41,6 +39,7 @@
     add_method,
     deserialize_and_fixup_type,
 )
+from mypy.semanal_shared import find_dataclass_transform_spec
 from mypy.server.trigger import make_wildcard_trigger
 from mypy.state import state
 from mypy.typeops import map_type_from_supertype
@@ -679,26 +678,20 @@ def _get_transform_spec(reason: Expression) -> DataclassTransformSpec:
     In those cases, we return a default spec rather than one based on a call to
     `typing.dataclass_transform`.
     """
-    node: Node | None = reason
+    if _is_dataclasses_decorator(reason):
+        return _TRANSFORM_SPEC_FOR_DATACLASSES
 
-    # The spec only lives on the function/class definition itself, so we need to unwrap down to that
-    # point
-    if isinstance(node, CallExpr):
-        # Decorators may take the form of either @decorator or @decorator(...); unwrap the latter
-        node = node.callee
-    if isinstance(node, RefExpr):
-        # If we see dataclasses.dataclass here, we know that we're not going to find a transform
-        # spec, so return early.
-        if node.fullname in dataclass_makers:
-            return _TRANSFORM_SPEC_FOR_DATACLASSES
-        node = node.node
-    if isinstance(node, Decorator):
-        node = node.func
-
-    if isinstance(node, SYMBOL_FUNCBASE_TYPES) and node.dataclass_transform_spec is not None:
-        return node.dataclass_transform_spec
-
-    raise AssertionError(
+    spec = find_dataclass_transform_spec(reason)
+    assert spec is not None, (
         "trying to find dataclass transform spec, but reason is neither dataclasses.dataclass nor "
         "decorated with typing.dataclass_transform"
     )
+    return spec
+
+
+def _is_dataclasses_decorator(node: Node) -> bool:
+    if isinstance(node, CallExpr):
+        node = node.callee
+    if isinstance(node, RefExpr):
+        return node.fullname in dataclass_makers
+    return False
diff --git a/mypy/semanal.py b/mypy/semanal.py
@@ -214,6 +214,7 @@
     PRIORITY_FALLBACKS,
     SemanticAnalyzerInterface,
     calculate_tuple_fallback,
+    find_dataclass_transform_spec,
     has_placeholder,
     set_callable_name as set_callable_name,
 )
@@ -1737,7 +1738,7 @@ def apply_class_plugin_hooks(self, defn: ClassDef) -> None:
                 # Special case: if the decorator is itself decorated with
                 # typing.dataclass_transform, apply the hook for the dataclasses plugin
                 # TODO: remove special casing here
-                if hook is None and is_dataclass_transform_decorator(decorator):
+                if hook is None and find_dataclass_transform_spec(decorator):
                     hook = dataclasses_plugin.dataclass_tag_callback
                 if hook:
                     hook(ClassDefContext(defn, decorator, self))
@@ -6687,24 +6688,3 @@ def halt(self, reason: str = ...) -> NoReturn:
     return isinstance(stmt, PassStmt) or (
         isinstance(stmt, ExpressionStmt) and isinstance(stmt.expr, EllipsisExpr)
     )
-
-
-def is_dataclass_transform_decorator(node: Node | None) -> bool:
-    if isinstance(node, RefExpr):
-        return is_dataclass_transform_decorator(node.node)
-    if isinstance(node, CallExpr):
-        # Like dataclasses.dataclass, transform-based decorators can be applied either with or
-        # without parameters; ie, both of these forms are accepted:
-        #
-        # @typing.dataclass_transform
-        # class Foo: ...
-        # @typing.dataclass_transform(eq=True, order=True, ...)
-        # class Bar: ...
-        #
-        # We need to unwrap the call for the second variant.
-        return is_dataclass_transform_decorator(node.callee)
-
-    if isinstance(node, Decorator) and node.func.dataclass_transform_spec is not None:
-        return True
-
-    return False
diff --git a/mypy/semanal_main.py b/mypy/semanal_main.py
@@ -41,7 +41,6 @@
 from mypy.semanal import (
     SemanticAnalyzer,
     apply_semantic_analyzer_patches,
-    is_dataclass_transform_decorator,
     remove_imported_names_from_symtable,
 )
 from mypy.semanal_classprop import (
@@ -51,6 +50,7 @@
     check_protocol_status,
 )
 from mypy.semanal_infer import infer_decorator_signature_if_simple
+from mypy.semanal_shared import find_dataclass_transform_spec
 from mypy.semanal_typeargs import TypeArgumentAnalyzer
 from mypy.server.aststrip import SavedAttributes
 from mypy.util import is_typeshed_file
@@ -467,7 +467,7 @@ def apply_hooks_to_class(
             # Special case: if the decorator is itself decorated with
             # typing.dataclass_transform, apply the hook for the dataclasses plugin
             # TODO: remove special casing here
-            if hook is None and is_dataclass_transform_decorator(decorator):
+            if hook is None and find_dataclass_transform_spec(decorator):
                 hook = dataclasses_plugin.dataclass_class_maker_callback
 
             if hook:

diff --git a/mypy/semanal_shared.py b/mypy/semanal_shared.py
@@ -11,10 +11,15 @@
 from mypy import join
 from mypy.errorcodes import ErrorCode
 from mypy.nodes import (
+    SYMBOL_FUNCBASE_TYPES,
+    CallExpr,
     Context,
+    DataclassTransformSpec,
+    Decorator,
     Expression,
     FuncDef,
     Node,
+    RefExpr,
     SymbolNode,
     SymbolTable,
     SymbolTableNode,
@@ -341,3 +346,40 @@ def visit_placeholder_type(self, t: PlaceholderType) -> bool:
 def has_placeholder(typ: Type) -> bool:
     """Check if a type contains any placeholder types (recursively)."""
     return typ.accept(HasPlaceholders())
+
+
+def find_dataclass_transform_spec(node: Node | None) -> DataclassTransformSpec | None:
+    """
+    Find the dataclass transform spec for the given node, if any exists.
+
+    Per PEP 681 (https://peps.python.org/pep-0681/#the-dataclass-transform-decorator), dataclass
+    transforms can be specified in multiple ways, including decorator functions and
+    metaclasses/base classes. This function resolves the spec from any of these variants.
+    """
+
+    # The spec only lives on the function/class definition itself, so we need to unwrap down to that
+    # point
+    if isinstance(node, CallExpr):
+        # Like dataclasses.dataclass, transform-based decorators can be applied either with or
+        # without parameters; ie, both of these forms are accepted:
+        #
+        # @typing.dataclass_transform
+        # class Foo: ...
+        # @typing.dataclass_transform(eq=True, order=True, ...)
+        # class Bar: ...
+        #
+        # We need to unwrap the call for the second variant.
+        node = node.callee
+
+    if isinstance(node, RefExpr):
+        node = node.node
+
+    if isinstance(node, Decorator):
+        # typing.dataclass_transform usage must always result in a Decorator; it always uses the
+        # `@dataclass_transform(...)` syntax and never `@dataclass_transform`
+        node = node.func
+
+    if isinstance(node, SYMBOL_FUNCBASE_TYPES):
+        return node.dataclass_transform_spec
+
+    return None