Handle code guarded by 'if TYPE_CHECKING' correctly

pyflakes/checker.py

@@ -321,10 +321,12 @@ class Binding(object):
                 the node that this binding was last used.
     """
 
-    def __init__(self, name, source):
+    def __init__(self, name, source, for_annotations):
         self.name = name
         self.source = source
         self.used = False
+        assert isinstance(for_annotations, bool)
+        self.for_annotations = for_annotations
 
     def __str__(self):
         return self.name
@@ -349,7 +351,7 @@ class Builtin(Definition):
     """A definition created for all Python builtins."""
 
     def __init__(self, name):
-        super(Builtin, self).__init__(name, None)
+        super(Builtin, self).__init__(name, None, for_annotations=False)
 
     def __repr__(self):
         return '<%s object %r at 0x%x>' % (self.__class__.__name__,
@@ -391,10 +393,11 @@ class Importation(Definition):
     @type fullName: C{str}
     """
 
-    def __init__(self, name, source, full_name=None):
+    def __init__(self, name, source, for_annotations, full_name=None):
         self.fullName = full_name or name
         self.redefined = []
-        super(Importation, self).__init__(name, source)
+        super(Importation, self).__init__(name, source,
+                                          for_annotations=for_annotations)
 
     def redefines(self, other):
         if isinstance(other, SubmoduleImportation):
@@ -439,11 +442,12 @@ class SubmoduleImportation(Importation):
     name is also the same, to avoid false positives.
     """
 
-    def __init__(self, name, source):
+    def __init__(self, name, source, for_annotations):
         # A dot should only appear in the name when it is a submodule import
         assert '.' in name and (not source or isinstance(source, ast.Import))
         package_name = name.split('.')[0]
-        super(SubmoduleImportation, self).__init__(package_name, source)
+        super(SubmoduleImportation, self).__init__(
+            package_name, source, for_annotations=for_annotations)
         self.fullName = name
 
     def redefines(self, other):
@@ -461,7 +465,7 @@ def source_statement(self):
 
 class ImportationFrom(Importation):
 
-    def __init__(self, name, source, module, real_name=None):
+    def __init__(self, name, source, module, for_annotations, real_name=None):
         self.module = module
         self.real_name = real_name or name
 
@@ -470,7 +474,8 @@ def __init__(self, name, source, module, real_name=None):
         else:
             full_name = module + '.' + self.real_name
 
-        super(ImportationFrom, self).__init__(name, source, full_name)
+        super(ImportationFrom, self).__init__(name, source, full_name=full_name,
+                                              for_annotations=for_annotations)
 
     def __str__(self):
         """Return import full name with alias."""
@@ -492,8 +497,9 @@ def source_statement(self):
 class StarImportation(Importation):
     """A binding created by a 'from x import *' statement."""
 
-    def __init__(self, name, source):
-        super(StarImportation, self).__init__('*', source)
+    def __init__(self, name, source, for_annotations):
+        super(StarImportation, self).__init__('*', source,
+                                              for_annotations=for_annotations)
         # Each star importation needs a unique name, and
         # may not be the module name otherwise it will be deemed imported
         self.name = name + '.*'
@@ -519,14 +525,17 @@ class FutureImportation(ImportationFrom):
     """
 
     def __init__(self, name, source, scope):
-        super(FutureImportation, self).__init__(name, source, '__future__')
+        super(FutureImportation, self).__init__(name, source, '__future__',
+                                                for_annotations=False)
         self.used = (scope, source)
 
 
 class Argument(Binding):
     """
     Represents binding a name as an argument.
     """
+    def __init__(self, name, source):
+        super(Argument, self).__init__(name, source, for_annotations=False)
 
 
 class Assignment(Binding):
@@ -576,7 +585,7 @@ class ExportBinding(Binding):
     C{__all__} will not have an unused import warning reported for them.
     """
 
-    def __init__(self, name, source, scope):
+    def __init__(self, name, source, scope, for_annotations):
         if '__all__' in scope and isinstance(source, ast.AugAssign):
             self.names = list(scope['__all__'].names)
         else:
@@ -607,7 +616,7 @@ def _add_to_names(container):
                 # If not list concatenation
                 else:
                     break
-        super(ExportBinding, self).__init__(name, source)
+        super(ExportBinding, self).__init__(name, source, for_annotations)
 
 
 class Scope(dict):
@@ -871,6 +880,7 @@ class Checker(object):
     traceTree = False
     _in_annotation = AnnotationState.NONE
     _in_deferred = False
+    _in_type_checking = False
 
     builtIns = set(builtin_vars).union(_MAGIC_GLOBALS)
     _customBuiltIns = os.environ.get('PYFLAKES_BUILTINS')
@@ -1205,12 +1215,16 @@ def handleNodeLoad(self, node):
                     self.report(messages.InvalidPrintSyntax, node)
 
             try:
-                scope[name].used = (self.scope, node)
+                n = scope[name]
+                if n.for_annotations and not self._in_annotation:
+                    # Only defined during type-checking; this does not count. Real code
+                    # (not an annotation) using this binding will not work.
+                    continue
+                n.used = (self.scope, node)
 
                 # if the name of SubImportation is same as
                 # alias of other Importation and the alias
                 # is used, SubImportation also should be marked as used.
-                n = scope[name]
                 if isinstance(n, Importation) and n._has_alias():
                     try:
                         scope[n.fullName].used = (self.scope, node)
@@ -1274,17 +1288,18 @@ def handleNodeStore(self, node):
 
         parent_stmt = self.getParent(node)
         if isinstance(parent_stmt, ANNASSIGN_TYPES) and parent_stmt.value is None:
-            binding = Annotation(name, node)
+            binding = Annotation(name, node, for_annotations=self._in_type_checking)
         elif isinstance(parent_stmt, (FOR_TYPES, ast.comprehension)) or (
                 parent_stmt != node._pyflakes_parent and
                 not self.isLiteralTupleUnpacking(parent_stmt)):
-            binding = Binding(name, node)
+            binding = Binding(name, node, for_annotations=self._in_type_checking)
         elif name == '__all__' and isinstance(self.scope, ModuleScope):
-            binding = ExportBinding(name, node._pyflakes_parent, self.scope)
+            binding = ExportBinding(name, node._pyflakes_parent, self.scope,
+                                    for_annotations=self._in_type_checking)
         elif PY2 and isinstance(getattr(node, 'ctx', None), ast.Param):
             binding = Argument(name, self.getScopeNode(node))
         else:
-            binding = Assignment(name, node)
+            binding = Assignment(name, node, for_annotations=self._in_type_checking)
         self.addBinding(node, binding)
 
     def handleNodeDelete(self, node):
@@ -1973,7 +1988,20 @@ def DICT(self, node):
     def IF(self, node):
         if isinstance(node.test, ast.Tuple) and node.test.elts != []:
             self.report(messages.IfTuple, node)
-        self.handleChildren(node)
+
+        prev = self._in_type_checking
+        if _is_typing(node.test, 'TYPE_CHECKING', self.scopeStack):
+            self._in_type_checking = True
+        # Handle the body of this if statement.
+        body_nodes = node.body
+        if not isinstance(body_nodes, list):
+            body_nodes = [body_nodes]
+        for body_node in body_nodes:
+            self.handleNode(body_node, node)
+        self._in_type_checking = prev
+
+        # Handle the test of the if statement (elif, else).
+        self.handleChildren(node, omit=["body"])
 
     IFEXP = IF
 
@@ -1994,7 +2022,8 @@ def GLOBAL(self, node):
 
             # One 'global' statement can bind multiple (comma-delimited) names.
             for node_name in node.names:
-                node_value = Assignment(node_name, node)
+                node_value = Assignment(node_name, node,
+                                        for_annotations=self._in_type_checking)
 
                 # Remove UndefinedName messages already reported for this name.
                 # TODO: if the global is not used in this scope, it does not
@@ -2096,7 +2125,8 @@ def FUNCTIONDEF(self, node):
         for deco in node.decorator_list:
             self.handleNode(deco, node)
         self.LAMBDA(node)
-        self.addBinding(node, FunctionDefinition(node.name, node))
+        self.addBinding(node, FunctionDefinition(
+            node.name, node, for_annotations=self._in_type_checking))
         # doctest does not process doctest within a doctest,
         # or in nested functions.
         if (self.withDoctest and
@@ -2221,7 +2251,8 @@ def CLASSDEF(self, node):
         for stmt in node.body:
             self.handleNode(stmt, node)
         self.popScope()
-        self.addBinding(node, ClassDefinition(node.name, node))
+        self.addBinding(node, ClassDefinition(
+            node.name, node, for_annotations=self._in_type_checking))
 
     def AUGASSIGN(self, node):
         self.handleNodeLoad(node.target)
@@ -2256,10 +2287,12 @@ def TUPLE(self, node):
     def IMPORT(self, node):
         for alias in node.names:
             if '.' in alias.name and not alias.asname:
-                importation = SubmoduleImportation(alias.name, node)
+                importation = SubmoduleImportation(
+                    alias.name, node, for_annotations=self._in_type_checking)
             else:
                 name = alias.asname or alias.name
-                importation = Importation(name, node, alias.name)
+                importation = Importation(name, node, full_name=alias.name,
+                                          for_annotations=self._in_type_checking)
             self.addBinding(node, importation)
 
     def IMPORTFROM(self, node):
@@ -2284,16 +2317,17 @@ def IMPORTFROM(self, node):
             elif alias.name == '*':
                 # Only Python 2, local import * is a SyntaxWarning
                 if not PY2 and not isinstance(self.scope, ModuleScope):
-                    self.report(messages.ImportStarNotPermitted,
-                                node, module)
+                    self.report(messages.ImportStarNotPermitted, node, module)
                     continue
 
                 self.scope.importStarred = True
                 self.report(messages.ImportStarUsed, node, module)
-                importation = StarImportation(module, node)
+                importation = StarImportation(
+                    module, node, for_annotations=self._in_type_checking)
             else:
-                importation = ImportationFrom(name, node,
-                                              module, alias.name)
+                importation = ImportationFrom(
+                    name, node, module, real_name=alias.name,
+                    for_annotations=self._in_type_checking)
             self.addBinding(node, importation)
 
     def TRY(self, node):

pyflakes/test/test_imports.py

@@ -14,78 +14,86 @@
 class TestImportationObject(TestCase):
 
     def test_import_basic(self):
-        binding = Importation('a', None, 'a')
+        binding = Importation('a', None, for_annotations=False, full_name='a')
         assert binding.source_statement == 'import a'
         assert str(binding) == 'a'
 
     def test_import_as(self):
-        binding = Importation('c', None, 'a')
+        binding = Importation('c', None, for_annotations=False, full_name='a')
         assert binding.source_statement == 'import a as c'
         assert str(binding) == 'a as c'
 
     def test_import_submodule(self):
-        binding = SubmoduleImportation('a.b', None)
+        binding = SubmoduleImportation('a.b', None, for_annotations=False)
         assert binding.source_statement == 'import a.b'
         assert str(binding) == 'a.b'
 
     def test_import_submodule_as(self):
         # A submodule import with an as clause is not a SubmoduleImportation
-        binding = Importation('c', None, 'a.b')
+        binding = Importation('c', None, for_annotations=False, full_name='a.b')
         assert binding.source_statement == 'import a.b as c'
         assert str(binding) == 'a.b as c'
 
     def test_import_submodule_as_source_name(self):
-        binding = Importation('a', None, 'a.b')
+        binding = Importation('a', None, for_annotations=False, full_name='a.b')
         assert binding.source_statement == 'import a.b as a'
         assert str(binding) == 'a.b as a'
 
     def test_importfrom_relative(self):
-        binding = ImportationFrom('a', None, '.', 'a')
+        binding = ImportationFrom('a', None, '.',
+                                  for_annotations=False, real_name='a')
         assert binding.source_statement == 'from . import a'
         assert str(binding) == '.a'
 
     def test_importfrom_relative_parent(self):
-        binding = ImportationFrom('a', None, '..', 'a')
+        binding = ImportationFrom('a', None, '..',
+                                  for_annotations=False, real_name='a')
         assert binding.source_statement == 'from .. import a'
         assert str(binding) == '..a'
 
     def test_importfrom_relative_with_module(self):
-        binding = ImportationFrom('b', None, '..a', 'b')
+        binding = ImportationFrom('b', None, '..a',
+                                  for_annotations=False, real_name='b')
         assert binding.source_statement == 'from ..a import b'
         assert str(binding) == '..a.b'
 
     def test_importfrom_relative_with_module_as(self):
-        binding = ImportationFrom('c', None, '..a', 'b')
+        binding = ImportationFrom('c', None, '..a',
+                                  for_annotations=False, real_name='b')
         assert binding.source_statement == 'from ..a import b as c'
         assert str(binding) == '..a.b as c'
 
     def test_importfrom_member(self):
-        binding = ImportationFrom('b', None, 'a', 'b')
+        binding = ImportationFrom('b', None, 'a',
+                                  for_annotations=False, real_name='b')
         assert binding.source_statement == 'from a import b'
         assert str(binding) == 'a.b'
 
     def test_importfrom_submodule_member(self):
-        binding = ImportationFrom('c', None, 'a.b', 'c')
+        binding = ImportationFrom('c', None, 'a.b',
+                                  for_annotations=False, real_name='c')
         assert binding.source_statement == 'from a.b import c'
         assert str(binding) == 'a.b.c'
 
     def test_importfrom_member_as(self):
-        binding = ImportationFrom('c', None, 'a', 'b')
+        binding = ImportationFrom('c', None, 'a',
+                                  for_annotations=False, real_name='b')
         assert binding.source_statement == 'from a import b as c'
         assert str(binding) == 'a.b as c'
 
     def test_importfrom_submodule_member_as(self):
-        binding = ImportationFrom('d', None, 'a.b', 'c')
+        binding = ImportationFrom('d', None, 'a.b',
+                                  for_annotations=False, real_name='c')
         assert binding.source_statement == 'from a.b import c as d'
         assert str(binding) == 'a.b.c as d'
 
     def test_importfrom_star(self):
-        binding = StarImportation('a.b', None)
+        binding = StarImportation('a.b', None, for_annotations=False)
         assert binding.source_statement == 'from a.b import *'
         assert str(binding) == 'a.b.*'
 
     def test_importfrom_star_relative(self):
-        binding = StarImportation('.b', None)
+        binding = StarImportation('.b', None, for_annotations=False)
         assert binding.source_statement == 'from .b import *'
         assert str(binding) == '.b.*'