[GR-31741] Github #178: __annotations__ attribute of a function have no "return" key

Author: ppisl

graalpython/com.oracle.graal.python.test/src/com/oracle/graal/python/test/parser/SSTSerializationTests.java

@@ -455,6 +455,13 @@ public void for01() throws Exception {
                                         "        return 10");
     }
 
+    @Test
+    public void functionAnnotationsTest() throws Exception {
+        checkSerialization("def a() -> int: pass");
+        checkSerialization("def a() -> max(1,3): pass");
+        checkSerialization("def fn(a:'x', b:int, c: list, d: 5+ 6) -> int: pass");
+    }
+
     @Test
     public void functionDecoratorTest() throws Exception {
         checkSerialization("@some.path.to.decorator\n" + "def fn(): pass");

graalpython/com.oracle.graal.python.test/src/tests/test_parser.py

@@ -690,7 +690,7 @@ def fn():
     assert "Function Documentation" in code1.co_consts
     assert "Function Documentation" not in code2.co_consts
     assert exec(code1) == exec(code2)
-
+    
 def test_annotations_in_global():
     test_globals = {'__annotations__': {}}
     code = compile ("a:int", "<test>", "exec")
@@ -725,7 +725,32 @@ def test_annotations_in_global():
     assert test_globals['__annotations__']['a'] == int
     assert test_globals['a'] == 1
 
+def test_annotations_in_function_declaration():
+    def fn1(a:int, b: 5+6): pass
+    assert len(fn1.__annotations__) == 2
+    assert fn1.__annotations__['a'] == int
+    assert fn1.__annotations__['b'] == 11
+    
+    def fn2(a: list, b: "Hello") -> int: pass
+    assert len(fn2.__annotations__) == 3
+    assert fn2.__annotations__['a'] == list
+    assert fn2.__annotations__['b'] == "Hello"
+    assert fn2.__annotations__['return'] == int
+
+    def fn3() -> sum((1,2,3,4)): pass
+    assert len(fn3.__annotations__) == 1
+    assert fn3.__annotations__['return'] == 10
+    
+    def fn4() -> f'hello {1+4}': pass
+    assert len(fn4.__annotations__) == 1
+    assert fn4.__annotations__['return'] == 'hello 5'
 
+    x = "Superman"
+    def fn5() -> f'hello {x}': pass
+    assert len(fn5.__annotations__) == 1
+    assert fn5.__annotations__['return'] == 'hello Superman'
+    
+
 def test_annotations_in_function():
     test_globals = {'__annotations__': {}}
     source = '''def fn():
@@ -878,4 +903,4 @@ def test_ComprehensionListExpr():
     assert [3,5,7] ==  [e+1 for e in [i*2 for i in (1,2,3)]]
     assert [('a', 1), ('a', 2), ('b', 1), ('b', 2)] == [ (c,s) for c in ('a','b') for s in (1,2)]
     assert [(1, 'a'), (2, 'a'), (1, 'b'), (2, 'b')] == [ (s,c) for c in ('a','b') for s in (1,2)]
-    
+    

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/parser/antlr/Python3.g4

@@ -563,10 +563,11 @@ decorated:
 
 async_funcdef: ASYNC funcdef;
 funcdef
-:
+:       
+        { SSTNode resultType = null; }
 	'def' n=NAME parameters
 	(
-		'->' test
+		'->' test { resultType = $test.result; }
 	)? ':' 
 	{ 
             String name = factory.mangleNameInCurrentScope($n.getText());
@@ -579,7 +580,7 @@ funcdef
         }
 	s = suite
 	{ 
-            SSTNode funcDef = new FunctionDefSSTNode(scopeEnvironment.getCurrentScope(), name, enclosingClassName, $parameters.result, optimize($s.result), getStartIndex(_localctx), getStopIndex(((FuncdefContext)_localctx).s));
+            SSTNode funcDef = new FunctionDefSSTNode(scopeEnvironment.getCurrentScope(), name, enclosingClassName, $parameters.result, optimize($s.result), resultType, getStartIndex(_localctx), getStopIndex(((FuncdefContext)_localctx).s));
             scopeEnvironment.popScope();
             loopState = savedLoopState;
             push(funcDef);