From 91773b89097927a2393bc5223295d37ed26e1632 Mon Sep 17 00:00:00 2001
From: =?utf8?q?Felix=20Hild=C3=A9n?= <felix.hilden@gmail.com>
Date: Tue, 13 Jul 2021 20:24:55 +0300
Subject: [PATCH] Improve AST safety parsing error message (#2304)

Co-authored-by: Hasan Ramezani <hasan.r67@gmail.com>
---
 CHANGES.md           |  1 +
 src/black/parsing.py | 52 +++++++++++++++++++++++++-------------------
 tests/test_black.py  | 32 ---------------------------
 3 files changed, 31 insertions(+), 54 deletions(-)

diff --git a/CHANGES.md b/CHANGES.md
index 611094f..b3224e1 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -20,6 +20,7 @@
 - Fixed option usage when using the `--code` flag (#2259)
 - Do not call `uvloop.install()` when _Black_ is used as a library (#2303)
 - Added `--required-version` option to require a specific version to be running (#2300)
+- Provide a more useful error when parsing fails during AST safety checks (#2304)
 - Fix incorrect custom breakpoint indices when string group contains fake f-strings
   (#2311)
 - Fix regression where `R` prefixes would be lowercased for docstrings (#2285)
diff --git a/src/black/parsing.py b/src/black/parsing.py
index 8e9feea..0b8d984 100644
--- a/src/black/parsing.py
+++ b/src/black/parsing.py
@@ -3,7 +3,7 @@ Parse Python code and perform AST validation.
 """
 import ast
 import sys
-from typing import Iterable, Iterator, List, Set, Union
+from typing import Iterable, Iterator, List, Set, Union, Tuple
 
 # lib2to3 fork
 from blib2to3.pytree import Node, Leaf
@@ -106,28 +106,36 @@ def lib2to3_unparse(node: Node) -> str:
     return code
 
 
-def parse_ast(src: str) -> Union[ast.AST, ast3.AST, ast27.AST]:
+def parse_single_version(
+    src: str, version: Tuple[int, int]
+) -> Union[ast.AST, ast3.AST, ast27.AST]:
     filename = "<unknown>"
-    if sys.version_info >= (3, 8):
-        # TODO: support Python 4+ ;)
-        for minor_version in range(sys.version_info[1], 4, -1):
-            try:
-                return ast.parse(src, filename, feature_version=(3, minor_version))
-            except SyntaxError:
-                continue
-    else:
-        for feature_version in (7, 6):
-            try:
-                return ast3.parse(src, filename, feature_version=feature_version)
-            except SyntaxError:
-                continue
-    if ast27.__name__ == "ast":
-        raise SyntaxError(
-            "The requested source code has invalid Python 3 syntax.\n"
-            "If you are trying to format Python 2 files please reinstall Black"
-            " with the 'python2' extra: `python3 -m pip install black[python2]`."
-        )
-    return ast27.parse(src)
+    # typed_ast is needed because of feature version limitations in the builtin ast
+    if sys.version_info >= (3, 8) and version >= (3,):
+        return ast.parse(src, filename, feature_version=version)
+    elif version >= (3,):
+        return ast3.parse(src, filename, feature_version=version[1])
+    elif version == (2, 7):
+        return ast27.parse(src)
+    raise AssertionError("INTERNAL ERROR: Tried parsing unsupported Python version!")
+
+
+def parse_ast(src: str) -> Union[ast.AST, ast3.AST, ast27.AST]:
+    # TODO: support Python 4+ ;)
+    versions = [(3, minor) for minor in range(3, sys.version_info[1] + 1)]
+
+    if ast27.__name__ != "ast":
+        versions.append((2, 7))
+
+    first_error = ""
+    for version in sorted(versions, reverse=True):
+        try:
+            return parse_single_version(src, version)
+        except SyntaxError as e:
+            if not first_error:
+                first_error = str(e)
+
+    raise SyntaxError(first_error)
 
 
 def stringify_ast(
diff --git a/tests/test_black.py b/tests/test_black.py
index 42ac119..e3be9c7 100644
--- a/tests/test_black.py
+++ b/tests/test_black.py
@@ -451,38 +451,6 @@ class BlackTestCase(BlackBaseTestCase):
             )
             self.assertEqual(expected, actual, msg)
 
-    @pytest.mark.no_python2
-    def test_python2_should_fail_without_optional_install(self) -> None:
-        if sys.version_info < (3, 8):
-            self.skipTest(
-                "Python 3.6 and 3.7 will install typed-ast to work and as such will be"
-                " able to parse Python 2 syntax without explicitly specifying the"
-                " python2 extra"
-            )
-
-        source = "x = 1234l"
-        tmp_file = Path(black.dump_to_file(source))
-        try:
-            runner = BlackRunner()
-            result = runner.invoke(black.main, [str(tmp_file)])
-            self.assertEqual(result.exit_code, 123)
-        finally:
-            os.unlink(tmp_file)
-        assert result.stderr_bytes is not None
-        actual = (
-            result.stderr_bytes.decode()
-            .replace("\n", "")
-            .replace("\\n", "")
-            .replace("\\r", "")
-            .replace("\r", "")
-        )
-        msg = (
-            "The requested source code has invalid Python 3 syntax."
-            "If you are trying to format Python 2 files please reinstall Black"
-            " with the 'python2' extra: `python3 -m pip install black[python2]`."
-        )
-        self.assertIn(msg, actual)
-
     @pytest.mark.python2
     @patch("black.dump_to_file", dump_to_stderr)
     def test_python2_print_function(self) -> None:
-- 
2.39.2