From 836acad86371578527408a4c8f968cde1302e130 Mon Sep 17 00:00:00 2001
From: Shantanu <12621235+hauntsaninja@users.noreply.github.com>
Date: Fri, 22 Mar 2024 19:13:53 -0700
Subject: [PATCH] Improve AST safety check (#4290)

Fixes #4288, regressed by #4270
---
 CHANGES.md           |  2 ++
 src/black/parsing.py |  6 ++----
 tests/test_black.py  | 16 ++++++++++++++++
 3 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/CHANGES.md b/CHANGES.md
index e6098b7ff59..014aec77392 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -10,6 +10,8 @@
 
 <!-- Changes that affect Black's stable style -->
 
+- Fix unwanted crashes caused by AST equivalency check (#4290)
+
 ### Preview style
 
 <!-- Changes that affect Black's preview style -->
diff --git a/src/black/parsing.py b/src/black/parsing.py
index aa97a8cecea..e8664b5ee23 100644
--- a/src/black/parsing.py
+++ b/src/black/parsing.py
@@ -223,11 +223,9 @@ def _stringify_ast(node: ast.AST, parent_stack: List[ast.AST]) -> Iterator[str]:
                 and field == "value"
                 and isinstance(value, str)
                 and len(parent_stack) >= 2
+                # Any standalone string, ideally this would
+                # exactly match black.nodes.is_docstring
                 and isinstance(parent_stack[-1], ast.Expr)
-                and isinstance(
-                    parent_stack[-2],
-                    (ast.FunctionDef, ast.AsyncFunctionDef, ast.Module, ast.ClassDef),
-                )
             ):
                 # Constant strings may be indented across newlines, if they are
                 # docstrings; fold spaces after newlines when comparing. Similarly,
diff --git a/tests/test_black.py b/tests/test_black.py
index 6ca67a287c8..2e3ae4503f5 100644
--- a/tests/test_black.py
+++ b/tests/test_black.py
@@ -2904,6 +2904,22 @@ async def f():
                 """docstring"""
             ''',
         )
+        self.check_ast_equivalence(
+            """
+            if __name__ == "__main__":
+                "  docstring-like  "
+            """,
+            '''
+            if __name__ == "__main__":
+                """docstring-like"""
+            ''',
+        )
+        self.check_ast_equivalence(r'def f(): r" \n "', r'def f(): "\\n"')
+        self.check_ast_equivalence('try: pass\nexcept: " x "', 'try: pass\nexcept: "x"')
+
+        self.check_ast_equivalence(
+            'def foo(): return " x "', 'def foo(): return "x"', should_fail=True
+        )
 
     def test_assert_equivalent_fstring(self) -> None:
         major, minor = sys.version_info[:2]