python · JukkaL · Apr 16, 2016 · Apr 7, 2016 · Apr 16, 2016 · Apr 16, 2016
diff --git a/mypy/semanal.py b/mypy/semanal.py
@@ -76,6 +76,7 @@
 from mypy.lex import lex
 from mypy.parsetype import parse_type
 from mypy.sametypes import is_same_type
+from mypy.erasetype import erase_typevars
 from mypy import defaults
 
 
@@ -2426,8 +2427,12 @@ def visit_class_def(self, tdef: ClassDef) -> None:
     def visit_decorator(self, dec: Decorator) -> None:
         """Try to infer the type of the decorated function.
 
-        This helps us resolve forward references to decorated
-        functions during type checking.
+        This lets us resolve references to decorated functions during
+        type checking when there are cyclic imports, as otherwise the
+        type might not be available when we need it.
+
+        This basically uses a simple special-purpose type inference
+        engine just for decorators.
         """
         super().visit_decorator(dec)
         if dec.var.is_property:
@@ -2453,13 +2458,21 @@ def visit_decorator(self, dec: Decorator) -> None:
                 decorator_preserves_type = False
                 break
         if decorator_preserves_type:
-            # No non-special decorators left. We can trivially infer the type
+            # No non-identity decorators left. We can trivially infer the type
             # of the function here.
             dec.var.type = function_type(dec.func, self.builtin_type('function'))
-        if dec.decorators and returns_any_if_called(dec.decorators[0]):
-            # The outermost decorator will return Any so we know the type of the
-            # decorated function.
-            dec.var.type = AnyType()
+        if dec.decorators:
+            if returns_any_if_called(dec.decorators[0]):
+                # The outermost decorator will return Any so we know the type of the
+                # decorated function.
+                dec.var.type = AnyType()
+            sig = find_fixed_callable_return(dec.decorators[0])
+            if sig:
+                # The outermost decorator always returns the same kind of function,
+                # so we know that this is the type of the decoratored function.
+                orig_sig = function_type(dec.func, self.builtin_type('function'))
+                sig.name = orig_sig.items()[0].name
+                dec.var.type = sig
 
     def visit_assignment_stmt(self, s: AssignmentStmt) -> None:
         self.analyze(s.type)
@@ -2673,3 +2686,23 @@ def returns_any_if_called(expr: Node) -> bool:
     elif isinstance(expr, CallExpr):
         return returns_any_if_called(expr.callee)
     return False
+
+
+def find_fixed_callable_return(expr: Node) -> Optional[CallableType]:
+    if isinstance(expr, RefExpr):
+        if isinstance(expr.node, FuncDef):
+            typ = expr.node.type
+            if typ:
+                if isinstance(typ, CallableType) and has_no_typevars(typ.ret_type):
+                    if isinstance(typ.ret_type, CallableType):
+                        return typ.ret_type
+    elif isinstance(expr, CallExpr):
+        t = find_fixed_callable_return(expr.callee)
+        if t:
+            if isinstance(t.ret_type, CallableType):
+                return t.ret_type
+    return None
+
+
+def has_no_typevars(typ: Type) -> bool:
+    return is_same_type(typ, erase_typevars(typ))
diff --git a/mypy/test/data/check-functions.test b/mypy/test/data/check-functions.test
@@ -863,6 +863,146 @@ def g(): pass
 def dec(f): pass
 
 
+-- Decorator functions in import cycles
+-- ------------------------------------
+
+
+[case testDecoratorWithIdentityTypeInImportCycle]
+import a
+
+[file a.py]
+import b
+from d import dec
+@dec
+def f(x: int) -> None: pass
+b.g(1) # E
+
+[file b.py]
+import a
+from d import dec
+@dec
+def g(x: str) -> None: pass
+a.f('')
+
+[file d.py]
+from typing import TypeVar
+T = TypeVar('T')
+def dec(f: T) -> T: return f
+
+[out]
+tmp/a.py:1: note: In module imported here,
+main:1: note: ... from here:
+tmp/b.py:5: error: Argument 1 to "f" has incompatible type "str"; expected "int"
+main:1: note: In module imported here:
+tmp/a.py:5: error: Argument 1 to "g" has incompatible type "int"; expected "str"
+
+[case testDecoratorWithNoAnnotationInImportCycle]
+import a
+
+[file a.py]
+import b
+from d import dec
+@dec
+def f(x: int) -> None: pass
+b.g(1, z=4)
+
+[file b.py]
+import a
+from d import dec
+@dec
+def g(x: str) -> None: pass
+a.f('', y=2)
+
+[file d.py]
+def dec(f): return f
+
+[case testDecoratorWithFixedReturnTypeInImportCycle]
+import a
+
+[file a.py]
+import b
+from d import dec
+@dec
+def f(x: int) -> str: pass
+b.g(1)()
+
+[file b.py]
+import a
+from d import dec
+@dec
+def g(x: int) -> str: pass
+a.f(1)()
+
+[file d.py]
+from typing import Callable
+def dec(f: Callable[[int], str]) -> Callable[[int], str]: return f
+
+[out]
+tmp/a.py:1: note: In module imported here,
+main:1: note: ... from here:
+tmp/b.py:5: error: "str" not callable
+main:1: note: In module imported here:
+tmp/a.py:5: error: "str" not callable
+
+[case testDecoratorWithCallAndFixedReturnTypeInImportCycle]
+import a
+
+[file a.py]
+import b
+from d import dec
+@dec()
+def f(x: int) -> str: pass
+b.g(1)()
+
+[file b.py]
+import a
+from d import dec
+@dec()
+def g(x: int) -> str: pass
+a.f(1)()
+
+[file d.py]
+from typing import Callable
+def dec() -> Callable[[Callable[[int], str]], Callable[[int], str]]: pass
+
+[out]
+tmp/a.py:1: note: In module imported here,
+main:1: note: ... from here:
+tmp/b.py:5: error: "str" not callable
+main:1: note: In module imported here:
+tmp/a.py:5: error: "str" not callable
+
+[case testDecoratorWithCallAndFixedReturnTypeInImportCycleAndDecoratorArgs]
+import a
+
+[file a.py]
+import b
+from d import dec
+@dec(1)
+def f(x: int) -> str: pass
+b.g(1)()
+
+[file b.py]
+import a
+from d import dec
+@dec(1)
+def g(x: int) -> str: pass
+a.f(1)()
+
+[file d.py]
+from typing import Callable
+def dec(x: str) -> Callable[[Callable[[int], str]], Callable[[int], str]]: pass
+
+[out]
+tmp/a.py:1: note: In module imported here,
+main:1: note: ... from here:
+tmp/b.py:3: error: Argument 1 to "dec" has incompatible type "int"; expected "str"
+tmp/b.py:5: error: "str" not callable
+main:1: note: In module imported here:
+tmp/a.py:3: error: Argument 1 to "dec" has incompatible type "int"; expected "str"
+tmp/a.py:5: error: "str" not callable
+
+
 -- Conditional function definition
 -- -------------------------------
 
@@ -1384,3 +1524,19 @@ with a:
     def f() -> None:
         pass
     f(1) # E: Too many arguments for "f"
+
+
+[case testNameForDecoratorMethod]
+from typing import Callable
+
+class A:
+    def f(self) -> None:
+        # In particular, test that the error message contains "g" of "A".
+        self.g() # E: Too few arguments for "g" of "A"
+        self.g(1)
+    @dec
+    def g(self, x: str) -> None: pass
+
+def dec(f: Callable[[A, str], None]) -> Callable[[A, int], None]: pass
+[out]
+main: note: In member "f" of class "A":
diff --git a/mypy/test/data/check-inference.test b/mypy/test/data/check-inference.test
@@ -1406,7 +1406,9 @@ x = 1
 [out]
 
 [case testMultipassAndDecoratedMethod]
-from typing import Callable
+from typing import Callable, TypeVar
+
+T = TypeVar('T')
 
 class A:
     def f(self) -> None:
@@ -1415,7 +1417,7 @@ class A:
     @dec
     def g(self, x: str) -> None: pass
 
-def dec(f: Callable[[A, str], None]) -> Callable[[A, int], None]: pass
+def dec(f: Callable[[A, str], T]) -> Callable[[A, int], T]: pass
 [out]
 main: note: In member "f" of class "A":
 

diff --git a/mypy/test/data/check-multiple-inheritance.test b/mypy/test/data/check-multiple-inheritance.test
@@ -233,7 +233,8 @@ class A(G[int]):
 class B(A, int): pass
 
 [case testCannotDetermineTypeInMultipleInheritance]
-from typing import Callable
+from typing import Callable, TypeVar
+T = TypeVar('T')
 class A(B, C):
     def f(self): pass
 class B:
@@ -242,8 +243,9 @@ class B:
 class C:
     @dec
     def f(self): pass
-def dec(f) -> Callable[[], None]: return None
+def dec(f: Callable[..., T]) -> Callable[..., T]:
+    return f
 [out]
 main: note: In class "A":
-main:2: error: Cannot determine type of 'f' in base class 'B'
-main:2: error: Cannot determine type of 'f' in base class 'C'
+main:3: error: Cannot determine type of 'f' in base class 'B'
+main:3: error: Cannot determine type of 'f' in base class 'C'