allow any expression

Instagram · May 22, 2023 · ff1a680 · ff1a680
1 parent a68ac53
commit ff1a680
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 12 deletions.
diff --git a/libcst/_nodes/statement.py b/libcst/_nodes/statement.py
@@ -21,7 +21,6 @@
     BaseAssignTargetExpression,
     BaseDelTargetExpression,
     BaseExpression,
-    Call,
     ConcatenatedString,
     ExpressionPosition,
     From,
@@ -1619,7 +1618,7 @@ class Decorator(CSTNode):
 
     #: The decorator that will return a new function wrapping the parent
     #: of this decorator.
-    decorator: Union[Name, Attribute, Call]
+    decorator: BaseExpression
 
     #: Line comments and empty lines before this decorator. The parent
     #: :class:`FunctionDef` or :class:`ClassDef` node owns leading lines before

diff --git a/libcst/_nodes/tests/test_funcdef.py b/libcst/_nodes/tests/test_funcdef.py
@@ -633,6 +633,20 @@ class FunctionDefCreationTest(CSTNodeTest):
                 ),
                 "code": "@bar()()\ndef foo(): pass\n",
             },
+            # Allow nested calls on decorator
+            {
+                "node": cst.FunctionDef(
+                    cst.Name("foo"),
+                    cst.Parameters(),
+                    cst.SimpleStatementSuite((cst.Pass(),)),
+                    (
+                        cst.Decorator(
+                            cst.BinaryOperation(cst.Name("a"), cst.Add(), cst.Name("b"))
+                        ),
+                    ),
+                ),
+                "code": "@a + b\ndef foo(): pass\n",
+            },
             # Allow parentheses around decorator
             {
                 "node": cst.FunctionDef(

diff --git a/libcst/matchers/__init__.py b/libcst/matchers/__init__.py
@@ -3089,13 +3089,6 @@ class Continue(BaseSmallStatement, BaseMatcherNode):
     ] = DoNotCare()
 
 
-NameOrAttributeOrCallMatchType = Union[
-    "Name",
-    "Attribute",
-    "Call",
-    MetadataMatchType,
-    MatchIfTrue[Union[cst.Name, cst.Attribute, cst.Call]],
-]
 TrailingWhitespaceMatchType = Union[
     "TrailingWhitespace", MetadataMatchType, MatchIfTrue[cst.TrailingWhitespace]
 ]
@@ -3104,10 +3097,10 @@ class Continue(BaseSmallStatement, BaseMatcherNode):
 @dataclass(frozen=True, eq=False, unsafe_hash=False)
 class Decorator(BaseMatcherNode):
     decorator: Union[
-        NameOrAttributeOrCallMatchType,
+        BaseExpressionMatchType,
         DoNotCareSentinel,
-        OneOf[NameOrAttributeOrCallMatchType],
-        AllOf[NameOrAttributeOrCallMatchType],
+        OneOf[BaseExpressionMatchType],
+        AllOf[BaseExpressionMatchType],
     ] = DoNotCare()
     leading_lines: Union[
         Sequence[