Specialize fused function local variables specified with pure-python (GH-3463)

These were previously getting missed. Added code to specialize them and tests to prove it.

Fixes https://github.com/cython/cython/issues/3142

Also fixes https://github.com/cython/cython/issues/3460 - (seems related enough to go in the same PR)

6 files changed, 30 insertions, 8 deletions

diff --git a/Cython/Compiler/ExprNodes.py b/Cython/Compiler/ExprNodes.py
index 1e9e4fcbb..0596904d3 100644
--- a/Cython/Compiler/ExprNodes.py
+++ b/Cython/Compiler/ExprNodes.py
@@ -1953,6 +1953,8 @@ class NameNode(AtomicExprNode):
             _, atype = annotation.analyse_type_annotation(env)
         if atype is None:
             atype = unspecified_type if as_target and env.directives['infer_types'] != False else py_object_type
+        if atype.is_fused and env.fused_to_specific:
+            atype = atype.specialize(env.fused_to_specific)
         self.entry = env.declare_var(name, atype, self.pos, is_cdef=not as_target)
         self.entry.annotation = annotation.expr
 
@@ -13593,10 +13595,16 @@ class AnnotationNode(ExprNode):
     # 2. The Cython use where the annotation can indicate an
     #  object type
     #
-    # doesn't handle the pre PEP-563 version where the
-    # annotation is evaluated into a Python Object
+    # Doesn't handle the pre PEP-563 version where the
+    # annotation is evaluated into a Python Object.
 
     subexprs = []
+
+    # 'untyped' is set for fused specializations:
+    # Once a fused function has been created we don't want
+    # annotations to override an already set type.
+    untyped = False
+
     def __init__(self, pos, expr, string=None):
         """string is expected to already be a StringNode or None"""
         ExprNode.__init__(self, pos)
@@ -13617,6 +13625,9 @@ class AnnotationNode(ExprNode):
         return self.analyse_type_annotation(env)[1]
 
     def analyse_type_annotation(self, env, assigned_value=None):
+        if self.untyped:
+            # Already applied as a fused type, not re-evaluating it here.
+            return None, None
         annotation = self.expr
         base_type = None
         is_ambiguous = False
diff --git a/Cython/Compiler/FusedNode.py b/Cython/Compiler/FusedNode.py
index 1ddb969ab..dbf3bbeac 100644
--- a/Cython/Compiler/FusedNode.py
+++ b/Cython/Compiler/FusedNode.py
@@ -220,6 +220,10 @@ class FusedCFuncDefNode(StatListNode):
                 arg.type = arg.type.specialize(fused_to_specific)
                 if arg.type.is_memoryviewslice:
                     arg.type.validate_memslice_dtype(arg.pos)
+                if arg.annotation:
+                    # TODO might be nice if annotations were specialized instead?
+                    # (Or might be hard to do reliably)
+                    arg.annotation.untyped = True
 
     def create_new_local_scope(self, node, env, f2s):
         """
diff --git a/Cython/Compiler/Nodes.py b/Cython/Compiler/Nodes.py
index 63948c067..e79f4a6ef 100644
--- a/Cython/Compiler/Nodes.py
+++ b/Cython/Compiler/Nodes.py
@@ -1680,8 +1680,6 @@ class FuncDefNode(StatNode, BlockNode):
             return arg
         if other_type is None:
             error(type_node.pos, "Not a type")
-        elif other_type.is_fused and any(orig_type.same_as(t) for t in other_type.types):
-            pass # use specialized rather than fused type
         elif orig_type is not py_object_type and not orig_type.same_as(other_type):
             error(arg.base_type.pos, "Signature does not agree with previous declaration")
             error(type_node.pos, "Previous declaration here")
diff --git a/Cython/Compiler/ParseTreeTransforms.py b/Cython/Compiler/ParseTreeTransforms.py
index 6d4830c25..608fe079b 100644
--- a/Cython/Compiler/ParseTreeTransforms.py
+++ b/Cython/Compiler/ParseTreeTransforms.py
@@ -1890,6 +1890,8 @@ if VALUE is not None:
         for var, type_node in node.directive_locals.items():
             if not lenv.lookup_here(var):   # don't redeclare args
                 type = type_node.analyse_as_type(lenv)
+                if type and type.is_fused and lenv.fused_to_specific:
+                    type = type.specialize(lenv.fused_to_specific)
                 if type:
                     lenv.declare_var(var, type, type_node.pos)
                 else:
diff --git a/tests/run/pure_fused.pxd b/tests/run/pure_fused.pxd
index 8517227d0..c00105087 100644
--- a/tests/run/pure_fused.pxd
+++ b/tests/run/pure_fused.pxd
@@ -1,6 +1,9 @@
+cimport cython
+
 ctypedef fused NotInPy:
     int
     float
 
 cdef class TestCls:
+    @cython.locals(loc = NotInPy)
     cpdef cpfunc(self, NotInPy arg)
diff --git a/tests/run/pure_fused.py b/tests/run/pure_fused.py
index b6d211e0f..35a9d27d7 100644
--- a/tests/run/pure_fused.py
+++ b/tests/run/pure_fused.py
@@ -1,5 +1,5 @@
 # mode: run
-# tag: fused, pure3.0
+# tag: fused, pure3.6
 
 #cython: annotation_typing=True
 
@@ -17,10 +17,11 @@ class TestCls:
         >>> TestCls().func1(2)
         'int'
         """
+        loc: 'NotInPy' = arg
         return cython.typeof(arg)
 
     if cython.compiled:
-        @cython.locals(arg = NotInPy)  # NameError in pure Python
+        @cython.locals(arg=NotInPy, loc=NotInPy)  # NameError for 'NotInPy' in pure Python
         def func2(self, arg):
             """
             >>> TestCls().func2(1.0)
@@ -28,6 +29,7 @@ class TestCls:
             >>> TestCls().func2(2)
             'int'
             """
+            loc = arg
             return cython.typeof(arg)
 
     def cpfunc(self, arg):
@@ -37,6 +39,7 @@ class TestCls:
         >>> TestCls().cpfunc(2)
         'int'
         """
+        loc = arg
         return cython.typeof(arg)
 
     def func1_inpy(self, arg: InPy):
@@ -46,9 +49,10 @@ class TestCls:
         >>> TestCls().func1_inpy(2)
         'int'
         """
+        loc: InPy = arg
         return cython.typeof(arg)
 
-    @cython.locals(arg = InPy)
+    @cython.locals(arg = InPy, loc = InPy)
     def func2_inpy(self, arg):
         """
         >>> TestCls().func2_inpy(1.0)
@@ -56,5 +60,5 @@ class TestCls:
         >>> TestCls().func2_inpy(2)
         'int'
         """
+        loc = arg
         return cython.typeof(arg)
-