From 94fe0bfa3fcf6ee7ecf94241dbf336d29ea52609 Mon Sep 17 00:00:00 2001
From: Jimmy Lai <yurinai@gmail.com>
Date: Thu, 10 Oct 2019 17:40:05 -0700
Subject: [PATCH] handle ComprehensionScope in QualifiedName

---
 libcst/metadata/scope_provider.py             |  2 +
 .../tests/test_qualified_name_provider.py     | 66 +++++++++++++++++--
 2 files changed, 63 insertions(+), 5 deletions(-)

diff --git a/libcst/metadata/scope_provider.py b/libcst/metadata/scope_provider.py
index b77bd2f1..486c57a1 100644
--- a/libcst/metadata/scope_provider.py
+++ b/libcst/metadata/scope_provider.py
@@ -281,6 +281,8 @@ class _NameUtil:
                 name_prefixes.append(f"{scope.name}.<locals>")
             elif isinstance(scope, GlobalScope):
                 break
+            elif isinstance(scope, ComprehensionScope):
+                name_prefixes.append("<comprehension>")
             else:
                 raise Exception(f"Unexpected Scope: {scope}")
             scope = scope.parent
diff --git a/libcst/metadata/tests/test_qualified_name_provider.py b/libcst/metadata/tests/test_qualified_name_provider.py
index 134d7df1..fd82910a 100644
--- a/libcst/metadata/tests/test_qualified_name_provider.py
+++ b/libcst/metadata/tests/test_qualified_name_provider.py
@@ -152,11 +152,11 @@ class ScopeProviderTest(UnitTest):
     def test_multiple_assignments(self) -> None:
         m, names = get_qualified_name_metadata_provider(
             """
-                if 1:
-                    from a import b as c
-                elif 2:
-                    from d import e as c
-                c()
+            if 1:
+                from a import b as c
+            elif 2:
+                from d import e as c
+            c()
             """
         )
         call = ensure_type(
@@ -169,3 +169,59 @@ class ScopeProviderTest(UnitTest):
                 QualifiedName(name="d.e", source=QualifiedNameSource.IMPORT),
             },
         )
+
+    def test_comprehension(self) -> None:
+        m, names = get_qualified_name_metadata_provider(
+            """
+            class C:
+                def fn(self) -> None:
+                    [[k for k in i] for i in [j for j in range(10)]]
+                    # Note:
+                    # The qualified name of i is straightforward to be "C.fn.<locals>.<comprehension>.i".
+                    # ListComp j is evaluated outside of the ListComp i.
+                    # so j has qualified name "C.fn.<locals>.<comprehension>.j".
+                    # ListComp k is evaluated inside ListComp i.
+                    # so k has qualified name "C.fn.<locals>.<comprehension>.<comprehension>.k".
+            """
+        )
+        cls_def = ensure_type(m.body[0], cst.ClassDef)
+        fn_def = ensure_type(cls_def.body.body[0], cst.FunctionDef)
+        outer_comp = ensure_type(
+            ensure_type(
+                ensure_type(fn_def.body.body[0], cst.SimpleStatementLine).body[0],
+                cst.Expr,
+            ).value,
+            cst.ListComp,
+        )
+        i = outer_comp.for_in.target
+        self.assertEqual(
+            names[i],
+            {
+                QualifiedName(
+                    name="C.fn.<locals>.<comprehension>.i",
+                    source=QualifiedNameSource.LOCAL,
+                )
+            },
+        )
+        inner_comp_j = ensure_type(outer_comp.for_in.iter, cst.ListComp)
+        j = inner_comp_j.for_in.target
+        self.assertEqual(
+            names[j],
+            {
+                QualifiedName(
+                    name="C.fn.<locals>.<comprehension>.j",
+                    source=QualifiedNameSource.LOCAL,
+                )
+            },
+        )
+        inner_comp_k = ensure_type(outer_comp.elt, cst.ListComp)
+        k = inner_comp_k.for_in.target
+        self.assertEqual(
+            names[k],
+            {
+                QualifiedName(
+                    name="C.fn.<locals>.<comprehension>.<comprehension>.k",
+                    source=QualifiedNameSource.LOCAL,
+                )
+            },
+        )