[Python-checkins] r79916 - in sandbox/trunk/2to3/lib2to3: refactor.py tests/test_refactor.py

[benjamin.peterson]

Author: benjamin.peterson
Date: Fri Apr  9 23:05:21 2010
New Revision: 79916

Log:
generalize detection of __future__ imports and attach them to the tree

Modified:
   sandbox/trunk/2to3/lib2to3/refactor.py
   sandbox/trunk/2to3/lib2to3/tests/test_refactor.py

Modified: sandbox/trunk/2to3/lib2to3/refactor.py
==============================================================================
--- sandbox/trunk/2to3/lib2to3/refactor.py	(original)
+++ sandbox/trunk/2to3/lib2to3/refactor.py	Fri Apr  9 23:05:21 2010
@@ -122,13 +122,14 @@
     _to_system_newlines = _identity
 
 
-def _detect_future_print(source):
+def _detect_future_features(source):
     have_docstring = False
     gen = tokenize.generate_tokens(StringIO.StringIO(source).readline)
     def advance():
         tok = next(gen)
         return tok[0], tok[1]
     ignore = frozenset((token.NEWLINE, tokenize.NL, token.COMMENT))
+    features = set()
     try:
         while True:
             tp, value = advance()
@@ -140,26 +141,25 @@
                 have_docstring = True
             elif tp == token.NAME and value == u"from":
                 tp, value = advance()
-                if tp != token.NAME and value != u"__future__":
+                if tp != token.NAME or value != u"__future__":
                     break
                 tp, value = advance()
-                if tp != token.NAME and value != u"import":
+                if tp != token.NAME or value != u"import":
                     break
                 tp, value = advance()
                 if tp == token.OP and value == u"(":
                     tp, value = advance()
                 while tp == token.NAME:
-                    if value == u"print_function":
-                        return True
+                    features.add(value)
                     tp, value = advance()
-                    if tp != token.OP and value != u",":
+                    if tp != token.OP or value != u",":
                         break
                     tp, value = advance()
             else:
                 break
     except StopIteration:
         pass
-    return False
+    return frozenset(features)
 
 
 class FixerError(Exception):
@@ -341,7 +341,8 @@
             An AST corresponding to the refactored input stream; None if
             there were errors during the parse.
         """
-        if _detect_future_print(data):
+        features = _detect_future_features(data)
+        if "print_function" in features:
             self.driver.grammar = pygram.python_grammar_no_print_statement
         try:
             tree = self.driver.parse_string(data)
@@ -351,6 +352,7 @@
             return
         finally:
             self.driver.grammar = self.grammar
+        tree.future_features = features
         self.log_debug("Refactoring %s", name)
         self.refactor_tree(tree, name)
         return tree

Modified: sandbox/trunk/2to3/lib2to3/tests/test_refactor.py
==============================================================================
--- sandbox/trunk/2to3/lib2to3/tests/test_refactor.py	(original)
+++ sandbox/trunk/2to3/lib2to3/tests/test_refactor.py	Fri Apr  9 23:05:21 2010
@@ -61,42 +61,50 @@
         self.assertEqual(full_names,
                          ["myfixes.fix_" + name for name in contents])
 
-    def test_detect_future_print(self):
-        run = refactor._detect_future_print
-        self.assertFalse(run(""))
-        self.assertTrue(run("from __future__ import print_function"))
-        self.assertFalse(run("from __future__ import generators"))
-        self.assertFalse(run("from __future__ import generators, feature"))
-        input = "from __future__ import generators, print_function"
-        self.assertTrue(run(input))
-        input ="from __future__ import print_function, generators"
-        self.assertTrue(run(input))
-        input = "from __future__ import (print_function,)"
-        self.assertTrue(run(input))
-        input = "from __future__ import (generators, print_function)"
-        self.assertTrue(run(input))
-        input = "from __future__ import (generators, nested_scopes)"
-        self.assertFalse(run(input))
-        input = """from __future__ import generators
+    def test_detect_future_features(self):
+        run = refactor._detect_future_features
+        fs = frozenset
+        empty = fs()
+        self.assertEqual(run(""), empty)
+        self.assertEqual(run("from __future__ import print_function"),
+                         fs(("print_function",)))
+        self.assertEqual(run("from __future__ import generators"),
+                         fs(("generators",)))
+        self.assertEqual(run("from __future__ import generators, feature"),
+                         fs(("generators", "feature")))
+        inp = "from __future__ import generators, print_function"
+        self.assertEqual(run(inp), fs(("generators", "print_function")))
+        inp ="from __future__ import print_function, generators"
+        self.assertEqual(run(inp), fs(("print_function", "generators")))
+        inp = "from __future__ import (print_function,)"
+        self.assertEqual(run(inp), fs(("print_function",)))
+        inp = "from __future__ import (generators, print_function)"
+        self.assertEqual(run(inp), fs(("generators", "print_function")))
+        inp = "from __future__ import (generators, nested_scopes)"
+        self.assertEqual(run(inp), fs(("generators", "nested_scopes")))
+        inp = """from __future__ import generators
 from __future__ import print_function"""
-        self.assertTrue(run(input))
-        self.assertFalse(run("from"))
-        self.assertFalse(run("from 4"))
-        self.assertFalse(run("from x"))
-        self.assertFalse(run("from x 5"))
-        self.assertFalse(run("from x im"))
-        self.assertFalse(run("from x import"))
-        self.assertFalse(run("from x import 4"))
-        input = "'docstring'\nfrom __future__ import print_function"
-        self.assertTrue(run(input))
-        input = "'docstring'\n'somng'\nfrom __future__ import print_function"
-        self.assertFalse(run(input))
-        input = "# comment\nfrom __future__ import print_function"
-        self.assertTrue(run(input))
-        input = "# comment\n'doc'\nfrom __future__ import print_function"
-        self.assertTrue(run(input))
-        input = "class x: pass\nfrom __future__ import print_function"
-        self.assertFalse(run(input))
+        self.assertEqual(run(inp), fs(("generators", "print_function")))
+        invalid = ("from",
+                   "from 4",
+                   "from x",
+                   "from x 5",
+                   "from x im",
+                   "from x import",
+                   "from x import 4",
+                   )
+        for inp in invalid:
+            self.assertEqual(run(inp), empty)
+        inp = "'docstring'\nfrom __future__ import print_function"
+        self.assertEqual(run(inp), fs(("print_function",)))
+        inp = "'docstring'\n'somng'\nfrom __future__ import print_function"
+        self.assertEqual(run(inp), empty)
+        inp = "# comment\nfrom __future__ import print_function"
+        self.assertEqual(run(inp), fs(("print_function",)))
+        inp = "# comment\n'doc'\nfrom __future__ import print_function"
+        self.assertEqual(run(inp), fs(("print_function",)))
+        inp = "class x: pass\nfrom __future__ import print_function"
+        self.assertEqual(run(inp), empty)
 
     def test_get_headnode_dict(self):
         class NoneFix(fixer_base.BaseFix):