Repository: cython/cython
Pull request: Specialize fused function local variables specified with pure-python
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)

Problem statement:
Bug pure-Python mode and fused types?
This .pyx file can be compiled:
```cython
import cython

import numpy as np
cimport numpy as np

ctypedef fused A:
    np.int_t[:]
    np.float_t[:]

cpdef func(A arg):
    cdef A arr = np.empty_like(arg)
    return arr
```

but this doesn't work whereas it seems to me that it should mean the same thing in pure-Python mode (Is it true?):

bug.py:
```python
import numpy as np

def func(arg):
    arr = np.empty_like(arg)
    return arr
```

bug.pxd:
```
import cython

import numpy as np
cimport numpy as np

ctypedef fused A:
    np.int_t[:]
    np.float_t[:]

@cython.locals(arr=A)
cpdef func(A arg)
```

The error log:

```
Compiling bug.py because it changed.
[1/1] Cythonizing bug.py

Error compiling Cython file:
------------------------------------------------------------
...
import numpy as np

def func(arg):
    arr = np.empty_like(arg)
                      ^
------------------------------------------------------------

bug.py:4:23: Cannot coerce to a type that is not specialized

Error compiling Cython file:
------------------------------------------------------------
...
import numpy as np

def func(arg):
    arr = np.empty_like(arg)
    return arr
          ^
------------------------------------------------------------

bug.py:5:11: Type is not specialized

Error compiling Cython file:
------------------------------------------------------------
...
import numpy as np

def func(arg):
    arr = np.empty_like(arg)
                      ^
------------------------------------------------------------

bug.py:4:23: Cannot coerce to a type that is not specialized

Error compiling Cython file:
------------------------------------------------------------
...
import numpy as np

def func(arg):
    arr = np.empty_like(arg)
    return arr
          ^
------------------------------------------------------------

bug.py:5:11: Type is not specialized

Error compiling Cython file:
------------------------------------------------------------
...
import numpy as np

def func(arg):
    arr = np.empty_like(arg)
   ^
------------------------------------------------------------

bug.py:4:4: Invalid use of fused types, type cannot be specialized

Error compiling Cython file:
------------------------------------------------------------
...
import numpy as np

def func(arg):
    arr = np.empty_like(arg)
   ^
------------------------------------------------------------

bug.py:4:4: Invalid use of fused types, type cannot be specialized
Traceback (most recent call last):
  File "/home/users/augier3pi/.pyenv/versions/3.7.3/lib/python3.7/runpy.py", line 193, in _run_module_as_main
    "__main__", mod_spec)
  File "/home/users/augier3pi/.pyenv/versions/3.7.3/lib/python3.7/runpy.py", line 85, in _run_code
    exec(code, run_globals)
  File "/home/users/augier3pi/Dev/transonic/transonic_cl/cythonize.py", line 11, in <module>
    ext_modules=cythonize(path, language_level=3), include_dirs=[np.get_include()]
  File "/home/users/augier3pi/.pyenv/versions/3.7.3/lib/python3.7/site-packages/Cython/Build/Dependencies.py", line 1096, in cythonize
    cythonize_one(*args)
  File "/home/users/augier3pi/.pyenv/versions/3.7.3/lib/python3.7/site-packages/Cython/Build/Dependencies.py", line 1219, in cythonize_one
    raise CompileError(None, pyx_file)
Cython.Compiler.Errors.CompileError: bug.py
```

Fused types ordering behaves different in pure Python mode
Cython normally attempts to resolve fused types in order. Therefore

```
cdef fused FusedType:
    SomeCdefClass
    object
```

will pick a specialization for `SomeCdefClass` first, even though `SomeCdefClass` is also an `object. In Pure Python mode this is lost somehow. See test-cases below.

Working .pyx version
```
import cython

cdef fused FT:
    C  # C should be used in preference
    object

cdef class C:
    pass

def f(FT x):
    """
    >>> f(C())  # C picked in preference to object
    C
    """
    print(cython.typeof(x))
```

pure version - BROKEN
```
import cython

# C appears first so should be used in preference
# "C" is a string so it's valid in Python. However, using non-string doesn't change results
FT = cython.fused_type("C", object)
FT2 = cython.fused_type("C", int)

@cython.cclass
class C:
    pass

@cython.locals(x = FT)
def f(x):
    """
    >>> f(C())  # FAILS - returns object
    C
    """
    print(cython.typeof(x))

@cython.locals(x = FT2)
def g(x):
    """
    >>> g(C())  # so C is recognised, just not in preference to object
    C
    """
    print(cython.typeof(x))
```

I don't really know where this is going wrong, but the arguments appear to be passed into the constructor of `PyrexTypes.FusedType` in the right order.

Hints from issue discussion:
The code looks reasonable to me, so I consider it a bug that it doesn't work. My guess is that the `directive_locals` handling in the `AnalyseDeclarationsTransform` does not take (enough?) care about fused types.

Further investigation and PR welcome.
This is looking like a bug I introduced in https://github.com/cython/cython/pull/3391/files#diff-4befc14b337ab1321460ed046735be7dR1684 .... working on how to fix it but the source is at least identified

Commit message: Add comment.

Patch to review:
diff --git a/Cython/Compiler/ExprNodes.py b/Cython/Compiler/ExprNodes.py
index 1e9e4fcbb31..0596904d3d9 100644
--- a/Cython/Compiler/ExprNodes.py
+++ b/Cython/Compiler/ExprNodes.py
@@ -1953,6 +1953,8 @@ def declare_from_annotation(self, env, as_target=False):
             _, 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 @@ def analyse_as_type(self, env):
         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 1ddb969abe8..dbf3bbeac50 100644
--- a/Cython/Compiler/FusedNode.py
+++ b/Cython/Compiler/FusedNode.py
@@ -220,6 +220,10 @@ def _specialize_function_args(self, args, fused_to_specific):
                 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 63948c06705..e79f4a6ef35 100644
--- a/Cython/Compiler/Nodes.py
+++ b/Cython/Compiler/Nodes.py
@@ -1680,8 +1680,6 @@ def align_argument_type(self, env, arg):
             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 6d4830c2574..608fe079b6a 100644
--- a/Cython/Compiler/ParseTreeTransforms.py
+++ b/Cython/Compiler/ParseTreeTransforms.py
@@ -1890,6 +1890,8 @@ def visit_FuncDefNode(self, node):
         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 8517227d036..c00105087b7 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 b6d211e0fb1..45b0018bccc 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,12 @@ def func1(self, arg: 'NotInPy'):
         >>> 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 in pure Python
         def func2(self, arg):
             """
             >>> TestCls().func2(1.0)
@@ -28,6 +30,7 @@ def func2(self, arg):
             >>> TestCls().func2(2)
             'int'
             """
+            loc = arg
             return cython.typeof(arg)
 
     def cpfunc(self, arg):
@@ -37,6 +40,7 @@ def cpfunc(self, arg):
         >>> TestCls().cpfunc(2)
         'int'
         """
+        loc = arg
         return cython.typeof(arg)
 
     def func1_inpy(self, arg: InPy):
@@ -46,9 +50,10 @@ def func1_inpy(self, arg: InPy):
         >>> 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 +61,6 @@ def func2_inpy(self, arg):
         >>> TestCls().func2_inpy(2)
         'int'
         """
+        loc = arg
         return cython.typeof(arg)
 


Code review comments:
diff_hunk: @@ -17,17 +17,20 @@ def func1(self, arg: 'NotInPy'):
         >>> 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 in pure Python
path: tests/run/pure_fused.py
text: @user2:
```suggestion
        @user1.locals(arg=NotInPy, loc=NotInPy)  # NameError for 'NotInPy' in pure Python
```