ScopeProvider: Record Access for Attributes and Decorators

libcst/metadata/scope_provider.py

@@ -324,7 +324,7 @@ def __iter__(self) -> Iterator[BaseAssignment]:
 
     def __getitem__(self, node: Union[str, cst.CSTNode]) -> Collection[BaseAssignment]:
         """Get assignments given a name str or :class:`~libcst.CSTNode` by ``scope.assignments[node]``"""
-        name = _NameUtil.get_name_for(node)
+        name = get_full_name_for_node(node)
         return set(self._assignments[name]) if name in self._assignments else set()
 
     def __contains__(self, node: Union[str, cst.CSTNode]) -> bool:
@@ -346,31 +346,14 @@ def __iter__(self) -> Iterator[Access]:
 
     def __getitem__(self, node: Union[str, cst.CSTNode]) -> Collection[Access]:
         """Get accesses given a name str or :class:`~libcst.CSTNode` by ``scope.accesses[node]``"""
-        name = _NameUtil.get_name_for(node)
+        name = get_full_name_for_node(node)
         return self._accesses[name] if name in self._accesses else set()
 
     def __contains__(self, node: Union[str, cst.CSTNode]) -> bool:
         """Check if a name str or :class:`~libcst.CSTNode` has any access by ``node in scope.accesses``"""
         return len(self[node]) > 0
 
 
-class _NameUtil:
-    @staticmethod
-    def get_name_for(node: Union[str, cst.CSTNode]) -> Optional[str]:
-        """A helper function to retrieve simple name str from a CSTNode or str"""
-        if isinstance(node, cst.Name):
-            return node.value
-        elif isinstance(node, str):
-            return node
-        elif isinstance(node, cst.Call):
-            return _NameUtil.get_name_for(node.func)
-        elif isinstance(node, cst.Subscript):
-            return _NameUtil.get_name_for(node.value)
-        elif isinstance(node, (cst.FunctionDef, cst.ClassDef)):
-            return _NameUtil.get_name_for(node.name)
-        return None
-
-
 class Scope(abc.ABC):
     """
     Base class of all scope classes. Scope object stores assignments from imports,

libcst/metadata/tests/test_scope_provider.py

@@ -251,6 +251,45 @@ def test_dotted_import_access(self) -> None:
         self.assertEqual(list(scope_of_module["x.y"])[0].references, set())
         self.assertEqual(scope_of_module.accesses["x.y"], set())
 
+    def test_dotted_import_access_reference_by_node(self) -> None:
+        m, scopes = get_scope_metadata_provider(
+            """
+            import a.b.c
+            a.b.c()
+            """
+        )
+        scope_of_module = scopes[m]
+        first_statement = ensure_type(m.body[1], cst.SimpleStatementLine)
+        call = ensure_type(
+            ensure_type(first_statement.body[0], cst.Expr).value, cst.Call
+        )
+
+        a_b_c_assignment = cast(ImportAssignment, list(scope_of_module["a.b.c"])[0])
+        a_b_c_access = list(a_b_c_assignment.references)[0]
+        self.assertEqual(scope_of_module.accesses[call], {a_b_c_access})
+        self.assertEqual(a_b_c_access.node, call.func)
+
+    def test_decorator_access_reference_by_node(self) -> None:
+        m, scopes = get_scope_metadata_provider(
+            """
+            import decorator
+
+            @decorator
+            def f():
+                pass
+            """
+        )
+        scope_of_module = scopes[m]
+        function_def = ensure_type(m.body[1], cst.FunctionDef)
+        decorator = function_def.decorators[0]
+        self.assertTrue("decorator" in scope_of_module)
+
+        decorator_assignment = cast(
+            ImportAssignment, list(scope_of_module["decorator"])[0]
+        )
+        decorator_access = list(decorator_assignment.references)[0]
+        self.assertEqual(scope_of_module.accesses[decorator], {decorator_access})
+
     def test_dotted_import_with_call_access(self) -> None:
         m, scopes = get_scope_metadata_provider(
             """