bpo-39702: Relax grammar restrictions on decorators (PEP 614) (GH-18570)

3 files changed, 89 insertions, 15 deletions

diff --git a/Lib/test/test_decorators.py b/Lib/test/test_decorators.py
index 8953f64..298979e 100644
--- a/Lib/test/test_decorators.py
+++ b/Lib/test/test_decorators.py
@@ -151,21 +151,18 @@ class TestDecorators(unittest.TestCase):
         self.assertEqual(counts['double'], 4)
 
     def test_errors(self):
-        # Test syntax restrictions - these are all compile-time errors:
-        #
-        for expr in [ "1+2", "x[3]", "(1, 2)" ]:
-            # Sanity check: is expr is a valid expression by itself?
-            compile(expr, "testexpr", "exec")
-
-            codestr = "@%s\ndef f(): pass" % expr
-            self.assertRaises(SyntaxError, compile, codestr, "test", "exec")
 
-        # You can't put multiple decorators on a single line:
-        #
-        self.assertRaises(SyntaxError, compile,
-                          "@f1 @f2\ndef f(): pass", "test", "exec")
+        # Test SyntaxErrors:
+        for stmt in ("x,", "x, y", "x = y", "pass", "import sys"):
+            compile(stmt, "test", "exec")  # Sanity check.
+            with self.assertRaises(SyntaxError):
+                compile(f"@{stmt}\ndef f(): pass", "test", "exec")
 
-        # Test runtime errors
+        # Test TypeErrors that used to be SyntaxErrors:
+        for expr in ("1.+2j", "[1, 2][-1]", "(1, 2)", "True", "...", "None"):
+            compile(expr, "test", "eval")  # Sanity check.
+            with self.assertRaises(TypeError):
+                exec(f"@{expr}\ndef f(): pass")
 
         def unimp(func):
             raise NotImplementedError
@@ -179,6 +176,13 @@ class TestDecorators(unittest.TestCase):
             code = compile(codestr, "test", "exec")
             self.assertRaises(exc, eval, code, context)
 
+    def test_expressions(self):
+        for expr in (
+            "(x,)", "(x, y)", "x := y", "(x := y)", "x @y", "(x @ y)", "x[0]",
+            "w[x].y.z", "w + x - (y + z)", "x(y)()(z)", "[w, x, y][z]", "x.y",
+        ):
+            compile(f"@{expr}\ndef f(): pass", "test", "exec")
+
     def test_double(self):
         class C(object):
             @funcattrs(abc=1, xyz="haha")
diff --git a/Lib/test/test_grammar.py b/Lib/test/test_grammar.py
index 2ed73e3..922a516 100644
--- a/Lib/test/test_grammar.py
+++ b/Lib/test/test_grammar.py
@@ -460,7 +460,7 @@ class GrammarTests(unittest.TestCase):
 
     def test_funcdef(self):
         ### [decorators] 'def' NAME parameters ['->' test] ':' suite
-        ### decorator: '@' dotted_name [ '(' [arglist] ')' ] NEWLINE
+        ### decorator: '@' namedexpr_test NEWLINE
         ### decorators: decorator+
         ### parameters: '(' [typedargslist] ')'
         ### typedargslist: ((tfpdef ['=' test] ',')*
@@ -666,6 +666,20 @@ class GrammarTests(unittest.TestCase):
         def f(x) -> list: pass
         self.assertEqual(f.__annotations__, {'return': list})
 
+        # Test expressions as decorators (PEP 614):
+        @False or null
+        def f(x): pass
+        @d := null
+        def f(x): pass
+        @lambda f: null(f)
+        def f(x): pass
+        @[..., null, ...][1]
+        def f(x): pass
+        @null(null)(null)
+        def f(x): pass
+        @[null][0].__call__.__call__
+        def f(x): pass
+
         # test closures with a variety of opargs
         closure = 1
         def f(): return closure
@@ -1515,13 +1529,27 @@ class GrammarTests(unittest.TestCase):
             def meth2(self, arg): pass
             def meth3(self, a1, a2): pass
 
-        # decorator: '@' dotted_name [ '(' [arglist] ')' ] NEWLINE
+        # decorator: '@' namedexpr_test NEWLINE
         # decorators: decorator+
         # decorated: decorators (classdef | funcdef)
         def class_decorator(x): return x
         @class_decorator
         class G: pass
 
+        # Test expressions as decorators (PEP 614):
+        @False or class_decorator
+        class H: pass
+        @d := class_decorator
+        class I: pass
+        @lambda c: class_decorator(c)
+        class J: pass
+        @[..., class_decorator, ...][1]
+        class K: pass
+        @class_decorator(class_decorator)(class_decorator)
+        class L: pass
+        @[class_decorator][0].__call__.__call__
+        class M: pass
+
     def test_dictcomps(self):
         # dictorsetmaker: ( (test ':' test (comp_for |
         #                                   (',' test ':' test)* [','])) |
diff --git a/Lib/test/test_parser.py b/Lib/test/test_parser.py
index 7295f66..73178f3 100644
--- a/Lib/test/test_parser.py
+++ b/Lib/test/test_parser.py
@@ -226,6 +226,27 @@ class RoundtripLegalSyntaxTestCase(unittest.TestCase):
         self.check_suite("@funcattrs()\n"
                          "def f(): pass")
 
+        self.check_suite("@False or x\n"
+                         "def f(): pass")
+        self.check_suite("@d := x\n"
+                         "def f(): pass")
+        self.check_suite("@lambda f: x(f)\n"
+                         "def f(): pass")
+        self.check_suite("@[..., x, ...][1]\n"
+                         "def f(): pass")
+        self.check_suite("@x(x)(x)\n"
+                         "def f(): pass")
+        self.check_suite("@(x, x)\n"
+                         "def f(): pass")
+        self.check_suite("@...\n"
+                         "def f(): pass")
+        self.check_suite("@None\n"
+                         "def f(): pass")
+        self.check_suite("@w @(x @y) @(z)\n"
+                         "def f(): pass")
+        self.check_suite("@w[x].y.z\n"
+                         "def f(): pass")
+
         # keyword-only arguments
         self.check_suite("def f(*, a): pass")
         self.check_suite("def f(*, a = 5): pass")
@@ -270,6 +291,27 @@ class RoundtripLegalSyntaxTestCase(unittest.TestCase):
                          "@decorator2\n"
                          "class foo():pass")
 
+        self.check_suite("@False or x\n"
+                         "class C: pass")
+        self.check_suite("@d := x\n"
+                         "class C: pass")
+        self.check_suite("@lambda f: x(f)\n"
+                         "class C: pass")
+        self.check_suite("@[..., x, ...][1]\n"
+                         "class C: pass")
+        self.check_suite("@x(x)(x)\n"
+                         "class C: pass")
+        self.check_suite("@(x, x)\n"
+                         "class C: pass")
+        self.check_suite("@...\n"
+                         "class C: pass")
+        self.check_suite("@None\n"
+                         "class C: pass")
+        self.check_suite("@w @(x @y) @(z)\n"
+                         "class C: pass")
+        self.check_suite("@w[x].y.z\n"
+                         "class C: pass")
+
     def test_import_from_statement(self):
         self.check_suite("from sys.path import *")
         self.check_suite("from sys.path import dirname")