Define is_name_token (and friends) to resolve some `type: ignore`s (GH-2714)

[etc/vim.git] / src / black / __init__.py

diff --git a/src/black/__init__.py b/src/black/__init__.py
index e2376c45617469d7890c0036c15fc4ac7139c27b..9bc8fc15c49d8136d7da05e10670dac4de793ec5 100644 (file)
--- a/src/black/__init__.py
+++ b/src/black/__init__.py
@@ -37,10 +37,11 @@ from mypy_extensions import mypyc_attr
 from black.const import DEFAULT_LINE_LENGTH, DEFAULT_INCLUDES, DEFAULT_EXCLUDES
 from black.const import STDIN_PLACEHOLDER
 from black.nodes import STARS, syms, is_simple_decorator_expression
+from black.nodes import is_string_token
 from black.lines import Line, EmptyLineTracker
 from black.linegen import transform_line, LineGenerator, LN
 from black.comments import normalize_fmt_off
-from black.mode import Mode, TargetVersion
+from black.mode import FUTURE_FLAG_TO_FEATURE, Mode, TargetVersion
 from black.mode import Feature, supports_feature, VERSION_TO_FEATURES
 from black.cache import read_cache, write_cache, get_cache_info, filter_cached, Cache
 from black.concurrency import cancel, shutdown, maybe_install_uvloop
@@ -1080,7 +1081,7 @@ def format_str(src_contents: str, *, mode: Mode) -> FileContent:
     if mode.target_versions:
         versions = mode.target_versions
     else:
-        versions = detect_target_versions(src_node)
+        versions = detect_target_versions(src_node, future_imports=future_imports)
 
     # TODO: fully drop support and this code hopefully in January 2022 :D
     if TargetVersion.PY27 in mode.target_versions or versions == {TargetVersion.PY27}:
@@ -1132,7 +1133,9 @@ def decode_bytes(src: bytes) -> Tuple[FileContent, Encoding, NewLine]:
         return tiow.read(), encoding, newline
 
 
-def get_features_used(node: Node) -> Set[Feature]:  # noqa: C901
+def get_features_used(  # noqa: C901
+    node: Node, *, future_imports: Optional[Set[str]] = None
+) -> Set[Feature]:
     """Return a set of (relatively) new Python features used in this file.
 
     Currently looking for:
@@ -1142,12 +1145,20 @@ def get_features_used(node: Node) -> Set[Feature]:  # noqa: C901
     - positional only arguments in function signatures and lambdas;
     - assignment expression;
     - relaxed decorator syntax;
+    - usage of __future__ flags (annotations);
     - print / exec statements;
     """
     features: Set[Feature] = set()
+    if future_imports:
+        features |= {
+            FUTURE_FLAG_TO_FEATURE[future_import]
+            for future_import in future_imports
+            if future_import in FUTURE_FLAG_TO_FEATURE
+        }
+
     for n in node.pre_order():
-        if n.type == token.STRING:
-            value_head = n.value[:2]  # type: ignore
+        if is_string_token(n):
+            value_head = n.value[:2]
             if value_head in {'f"', 'F"', "f'", "F'", "rf", "fr", "RF", "FR"}:
                 features.add(Feature.F_STRINGS)
 
@@ -1200,6 +1211,21 @@ def get_features_used(node: Node) -> Set[Feature]:  # noqa: C901
                         if argch.type in STARS:
                             features.add(feature)
 
+        elif (
+            n.type in {syms.return_stmt, syms.yield_expr}
+            and len(n.children) >= 2
+            and n.children[1].type == syms.testlist_star_expr
+            and any(child.type == syms.star_expr for child in n.children[1].children)
+        ):
+            features.add(Feature.UNPACKING_ON_FLOW)
+
+        elif (
+            n.type == syms.annassign
+            and len(n.children) >= 4
+            and n.children[3].type == syms.testlist_star_expr
+        ):
+            features.add(Feature.ANN_ASSIGN_EXTENDED_RHS)
+
         # Python 2 only features (for its deprecation) except for integers, see above
         elif n.type == syms.print_stmt:
             features.add(Feature.PRINT_STMT)
@@ -1229,9 +1255,11 @@ def get_features_used(node: Node) -> Set[Feature]:  # noqa: C901
     return features
 
 
-def detect_target_versions(node: Node) -> Set[TargetVersion]:
+def detect_target_versions(
+    node: Node, *, future_imports: Optional[Set[str]] = None
+) -> Set[TargetVersion]:
     """Detect the version to target based on the nodes used."""
-    features = get_features_used(node)
+    features = get_features_used(node, future_imports=future_imports)
     return {
         version for version in TargetVersion if features <= VERSION_TO_FEATURES[version]
     }
@@ -1293,7 +1321,7 @@ def assert_equivalent(src: str, dst: str, *, pass_num: int = 1) -> None:
         src_ast = parse_ast(src)
     except Exception as exc:
         raise AssertionError(
-            "cannot use --safe with this file; failed to parse source file."
+            f"cannot use --safe with this file; failed to parse source file: {exc}"
         ) from exc
 
     try: