Add detecting future imports to config.

Several of the python 2 features are gated on these in addition to
version (like `with_statement`), and a refactoring tool like Bowler
commonly needs this information anyway.

libcst/_parser/detect_config.py

@@ -10,7 +10,7 @@
 from dataclasses import dataclass
 from io import BytesIO
 from tokenize import detect_encoding as py_tokenize_detect_encoding
-from typing import Iterable, Iterator, Pattern, Union
+from typing import Iterable, Iterator, Pattern, Set, Union
 
 from libcst._nodes.whitespace import NEWLINE_RE
 from libcst._parser.parso.python.token import PythonTokenTypes, TokenType
@@ -21,6 +21,10 @@
 
 
 _INDENT: TokenType = PythonTokenTypes.INDENT
+_NAME: TokenType = PythonTokenTypes.NAME
+_NEWLINE: TokenType = PythonTokenTypes.NEWLINE
+_STRING: TokenType = PythonTokenTypes.STRING
+
 _FALLBACK_DEFAULT_NEWLINE = "\n"
 _FALLBACK_DEFAULT_INDENT = "    "
 _CONTINUATION_RE: Pattern[str] = re.compile(r"\\(\r\n?|\n)", re.UNICODE)
@@ -80,6 +84,38 @@ def _detect_trailing_newline(source_str: str) -> bool:
     )
 
 
+def _detect_future_imports(tokens: Iterable[Token]) -> Set[str]:
+    """
+    Finds __future__ imports in their proper locations.
+
+    See `https://www.python.org/dev/peps/pep-0236/`_
+    """
+    future_imports: Set[str] = set()
+    state = 0
+    for tok in tokens:
+        if state == 0 and tok.type in (_STRING, _NEWLINE):
+            continue
+        elif state == 0 and tok.string == "from":
+            state = 1
+        elif state == 1 and tok.string == "__future__":
+            state = 2
+        elif state == 2 and tok.string == "import":
+            state = 3
+        elif state == 3 and tok.string == "as":
+            state = 4
+        elif state == 3 and tok.type == _NAME:
+            future_imports.add(tok.string)
+        elif state == 4 and tok.type == _NAME:
+            state = 3
+        elif state == 3 and tok.string in "(),":
+            continue
+        elif state == 3 and tok.type == _NEWLINE:
+            state = 0
+        else:
+            break
+    return future_imports
+
+
 def detect_config(
     source: Union[str, bytes],
     *,
@@ -144,6 +180,14 @@ def detect_config(
     else:
         default_indent = partial_default_indent
 
+    partial_future_imports = partial.future_imports
+    if isinstance(partial_future_imports, AutoConfig):
+        # Same note as above re itertools.tee, we will consume tokens.
+        tokens, tokens_dup = itertools.tee(tokens)
+        future_imports = _detect_future_imports(tokens_dup)
+    else:
+        future_imports = partial_future_imports
+
     return ConfigDetectionResult(
         config=ParserConfig(
             lines=lines,
@@ -152,6 +196,7 @@ def detect_config(
             default_newline=default_newline,
             has_trailing_newline=has_trailing_newline,
             version=python_version,
+            future_imports=future_imports,
         ),
         tokens=tokens,
     )

libcst/_parser/tests/test_detect_config.py

@@ -4,6 +4,7 @@
 # LICENSE file in the root directory of this source tree.
 
 # pyre-strict
+import dataclasses
 from typing import Union
 
 from libcst._parser.detect_config import detect_config
@@ -27,6 +28,7 @@ class TestDetectConfig(UnitTest):
                     default_newline="\n",
                     has_trailing_newline=False,
                     version=PythonVersionInfo(3, 7),
+                    future_imports=set(),
                 ),
             },
             "detect_trailing_newline_disabled": {
@@ -41,6 +43,7 @@ class TestDetectConfig(UnitTest):
                     default_newline="\n",
                     has_trailing_newline=False,
                     version=PythonVersionInfo(3, 7),
+                    future_imports=set(),
                 ),
             },
             "detect_default_newline_disabled": {
@@ -55,6 +58,7 @@ class TestDetectConfig(UnitTest):
                     default_newline="\n",
                     has_trailing_newline=False,
                     version=PythonVersionInfo(3, 7),
+                    future_imports=set(),
                 ),
             },
             "newline_inferred": {
@@ -69,6 +73,7 @@ class TestDetectConfig(UnitTest):
                     default_newline="\r\n",
                     has_trailing_newline=True,
                     version=PythonVersionInfo(3, 7),
+                    future_imports=set(),
                 ),
             },
             "newline_partial_given": {
@@ -85,6 +90,7 @@ class TestDetectConfig(UnitTest):
                     default_newline="\n",  # The given partial disables inference
                     has_trailing_newline=True,
                     version=PythonVersionInfo(3, 7),
+                    future_imports=set(),
                 ),
             },
             "indent_inferred": {
@@ -99,6 +105,7 @@ class TestDetectConfig(UnitTest):
                     default_newline="\n",
                     has_trailing_newline=True,
                     version=PythonVersionInfo(3, 7),
+                    future_imports=set(),
                 ),
             },
             "indent_partial_given": {
@@ -115,6 +122,7 @@ class TestDetectConfig(UnitTest):
                     default_newline="\n",
                     has_trailing_newline=True,
                     version=PythonVersionInfo(3, 7),
+                    future_imports=set(),
                 ),
             },
             "encoding_inferred": {
@@ -134,6 +142,7 @@ class TestDetectConfig(UnitTest):
                     default_newline="\n",
                     has_trailing_newline=True,
                     version=PythonVersionInfo(3, 7),
+                    future_imports=set(),
                 ),
             },
             "encoding_partial_given": {
@@ -155,6 +164,7 @@ class TestDetectConfig(UnitTest):
                     default_newline="\n",
                     has_trailing_newline=True,
                     version=PythonVersionInfo(3, 7),
+                    future_imports=set(),
                 ),
             },
             "encoding_str_not_bytes_disables_inference": {
@@ -174,6 +184,7 @@ class TestDetectConfig(UnitTest):
                     default_newline="\n",
                     has_trailing_newline=True,
                     version=PythonVersionInfo(3, 7),
+                    future_imports=set(),
                 ),
             },
             "encoding_non_ascii_compatible_utf_16_with_bom": {
@@ -188,6 +199,7 @@ class TestDetectConfig(UnitTest):
                     default_newline="\n",
                     has_trailing_newline=False,
                     version=PythonVersionInfo(3, 7),
+                    future_imports=set(),
                 ),
             },
             "detect_trailing_newline_missing_newline": {
@@ -202,6 +214,7 @@ class TestDetectConfig(UnitTest):
                     default_newline="\n",
                     has_trailing_newline=False,
                     version=PythonVersionInfo(3, 7),
+                    future_imports=set(),
                 ),
             },
             "detect_trailing_newline_has_newline": {
@@ -216,6 +229,7 @@ class TestDetectConfig(UnitTest):
                     default_newline="\n",
                     has_trailing_newline=True,
                     version=PythonVersionInfo(3, 7),
+                    future_imports=set(),
                 ),
             },
             "detect_trailing_newline_missing_newline_after_line_continuation": {
@@ -230,6 +244,7 @@ class TestDetectConfig(UnitTest):
                     default_newline="\n",
                     has_trailing_newline=False,
                     version=PythonVersionInfo(3, 7),
+                    future_imports=set(),
                 ),
             },
             "detect_trailing_newline_has_newline_after_line_continuation": {
@@ -244,6 +259,50 @@ class TestDetectConfig(UnitTest):
                     default_newline="\n",
                     has_trailing_newline=True,
                     version=PythonVersionInfo(3, 7),
+                    future_imports=set(),
+                ),
+            },
+            "future_imports_in_correct_position": {
+                "source": b"# C\n''' D '''\nfrom __future__ import a as b\n",
+                "partial": PartialParserConfig(python_version="3.7"),
+                "detect_trailing_newline": True,
+                "detect_default_newline": True,
+                "expected_config": ParserConfig(
+                    lines=[
+                        "# C\n",
+                        "''' D '''\n",
+                        "from __future__ import a as b\n",
+                        "",
+                    ],
+                    encoding="utf-8",
+                    default_indent="    ",
+                    default_newline="\n",
+                    has_trailing_newline=True,
+                    version=PythonVersionInfo(3, 7),
+                    future_imports={"a"},
+                ),
+            },
+            "future_imports_in_mixed_position": {
+                "source": (
+                    b"from __future__ import a, b\nimport os\n"
+                    b"from __future__ import c\n"
+                ),
+                "partial": PartialParserConfig(python_version="3.7"),
+                "detect_trailing_newline": True,
+                "detect_default_newline": True,
+                "expected_config": ParserConfig(
+                    lines=[
+                        "from __future__ import a, b\n",
+                        "import os\n",
+                        "from __future__ import c\n",
+                        "",
+                    ],
+                    encoding="utf-8",
+                    default_indent="    ",
+                    default_newline="\n",
+                    has_trailing_newline=True,
+                    version=PythonVersionInfo(3, 7),
+                    future_imports={"a", "b"},
                 ),
             },
         }
@@ -258,11 +317,13 @@ def test_detect_module_config(
         expected_config: ParserConfig,
     ) -> None:
         self.assertEqual(
-            detect_config(
-                source,
-                partial=partial,
-                detect_trailing_newline=detect_trailing_newline,
-                detect_default_newline=detect_default_newline,
-            ).config,
-            expected_config,
+            dataclasses.asdict(
+                detect_config(
+                    source,
+                    partial=partial,
+                    detect_trailing_newline=detect_trailing_newline,
+                    detect_default_newline=detect_default_newline,
+                ).config
+            ),
+            dataclasses.asdict(expected_config),
         )

libcst/_parser/types/config.py

@@ -10,7 +10,7 @@
 import re
 from dataclasses import dataclass, field, fields
 from enum import Enum
-from typing import List, Pattern, Sequence, Union
+from typing import List, Pattern, Sequence, Set, Union
 
 from libcst._add_slots import add_slots
 from libcst._nodes.whitespace import NEWLINE_RE
@@ -45,6 +45,7 @@ class ParserConfig(BaseWhitespaceParserConfig):
     default_newline: str
     has_trailing_newline: bool
     version: PythonVersionInfo
+    future_imports: Set[str]
 
 
 class AutoConfig(Enum):
@@ -96,6 +97,9 @@ class PartialParserConfig:
     #: this value defaults to ``"utf-8"``.
     encoding: Union[str, AutoConfig] = AutoConfig.token
 
+    #: Detected ``__future__`` import names
+    future_imports: Union[Set[str], AutoConfig] = AutoConfig.token
+
     #: The indentation of the file, expressed as a series of tabs and/or spaces. This
     #: value is inferred from the contents of the parsed source code by default.
     default_indent: Union[str, AutoConfig] = AutoConfig.token