Allow __future__ in doctest

Replaces plain attribute Checker.futuresAllowed with a property
that supports __future__ in both module and doctest.

2 files changed, 32 insertions, 12 deletions

diff --git a/pyflakes/checker.py b/pyflakes/checker.py
index f8c8df8..719f6b4 100644
--- a/pyflakes/checker.py
+++ b/pyflakes/checker.py
@@ -158,6 +158,18 @@ class StarImportation(Importation):
         self.fullName = name
 
 
+class FutureImportation(Importation):
+    """
+    A binding created by a from `__future__` import statement.
+
+    `__future__` imports are implicitly used.
+    """
+
+    def __init__(self, name, source, scope):
+        super(FutureImportation, self).__init__(name, source)
+        self.used = (scope, source)
+
+
 class Argument(Binding):
     """
     Represents binding a name as an argument.
@@ -255,6 +267,7 @@ class GeneratorScope(Scope):
 
 class ModuleScope(Scope):
     """Scope for a module."""
+    _futures_allowed = True
 
 
 class DoctestScope(ModuleScope):
@@ -310,7 +323,6 @@ class Checker(object):
         self.withDoctest = withDoctest
         self.scopeStack = [ModuleScope()]
         self.exceptHandlers = [()]
-        self.futuresAllowed = True
         self.root = tree
         self.handleChildren(tree)
         self.runDeferred(self._deferredFunctions)
@@ -357,6 +369,20 @@ class Checker(object):
                 isinstance(self.scopeStack[1], DoctestScope))
 
     @property
+    def futuresAllowed(self):
+        if not all(isinstance(scope, ModuleScope)
+                   for scope in self.scopeStack):
+            return False
+
+        return self.scope._futures_allowed
+
+    @futuresAllowed.setter
+    def futuresAllowed(self, value):
+        assert value is False
+        if isinstance(self.scope, ModuleScope):
+            self.scope._futures_allowed = False
+
+    @property
     def scope(self):
         return self.scopeStack[-1]
 
@@ -1025,7 +1051,9 @@ class Checker(object):
 
         for alias in node.names:
             name = alias.asname or alias.name
-            if alias.name == '*':
+            if node.module == '__future__':
+                importation = FutureImportation(name, node, self.scope)
+            elif alias.name == '*':
                 # Only Python 2, local import * is a SyntaxWarning
                 if not PY2 and not isinstance(self.scope, ModuleScope):
                     self.report(messages.ImportStarNotPermitted,
@@ -1037,8 +1065,6 @@ class Checker(object):
                 importation = StarImportation(node.module, node)
             else:
                 importation = Importation(name, node)
-            if node.module == '__future__':
-                importation.used = (self.scope, node)
             self.addBinding(node, importation)
 
     def TRY(self, node):
diff --git a/pyflakes/test/test_doctests.py b/pyflakes/test/test_doctests.py
index 9dd4278..aed7957 100644
--- a/pyflakes/test/test_doctests.py
+++ b/pyflakes/test/test_doctests.py
@@ -431,18 +431,12 @@ class Test(TestCase):
 
 
 class TestOther(_DoctestMixin, TestOther):
-    pass
+    """Run TestOther with each test wrapped in a doctest."""
 
 
 class TestImports(_DoctestMixin, TestImports):
-
-    def test_futureImport(self):
-        """XXX This test can't work in a doctest"""
-
-    def test_futureImportUsed(self):
-        """XXX This test can't work in a doctest"""
+    """Run TestImports with each test wrapped in a doctest."""
 
 
 class TestUndefinedNames(_DoctestMixin, TestUndefinedNames):
     """Run TestUndefinedNames with each test wrapped in a doctest."""
-    pass