From 80127e4d0254ff6577adf6c04ca132a3f0a63278 Mon Sep 17 00:00:00 2001
From: Juan Altmayer Pizzorno <juan@altmayer.com>
Date: Wed, 21 Aug 2024 17:37:47 -0400
Subject: [PATCH] - added (simplistic) support for looking up assignments such
 as those   used for constants;

---
 src/coverup/codeinfo.py |  3 +++
 tests/test_codeinfo.py  | 37 +++++++++++++++++++++++++++++++++++++
 2 files changed, 40 insertions(+)

diff --git a/src/coverup/codeinfo.py b/src/coverup/codeinfo.py
index 9863126..dcfebaa 100644
--- a/src/coverup/codeinfo.py
+++ b/src/coverup/codeinfo.py
@@ -79,6 +79,9 @@ def _find_name_path(node: ast.AST, name: T.List[str]) -> T.List[ast.AST]:
         if isinstance(c, (ast.FunctionDef, ast.AsyncFunctionDef)) and c.name == name[0]:
             return [c] if len(name) == 1 else []
 
+        if isinstance(c, ast.Assign) and any(t.id == name[0] for t in c.targets):
+            return [c] if len(name) == 1 else []
+
         if isinstance(c, (ast.Import, ast.ImportFrom)):
             # FIXME the first matching import needn't be the one that resolves the name:
             #   import foo.bar
diff --git a/tests/test_codeinfo.py b/tests/test_codeinfo.py
index 153b7ab..b123449 100644
--- a/tests/test_codeinfo.py
+++ b/tests/test_codeinfo.py
@@ -250,6 +250,43 @@ def foo():
     assert codeinfo.get_info(tree, 'foo') == None
 
 
+def test_get_info_assignment():
+    code = textwrap.dedent("""\
+        PI = 3.1415
+
+        x = 0
+        if x == 0:
+            TYPES = (
+                int,
+                str,
+            )
+
+        class C:
+            x = 10
+
+            def __init__(self, x: int) -> C:
+                self._foo = x
+        """
+    )
+
+    tree = ast.parse(code)
+    tree.path = Path("foo.py")
+
+    assert codeinfo.get_info(tree, 'PI') == textwrap.dedent("""\
+        PI = 3.1415"""
+    )
+
+    # FIXME do proper code slicing
+    assert codeinfo.get_info(tree, 'TYPES') == textwrap.dedent("""\
+        TYPES = (int, str)"""
+    )
+
+    assert codeinfo.get_info(tree, 'C.x') == textwrap.dedent("""\
+        class C:
+            ...
+            x = 10"""
+    )
+
 def test_get_info_imported(import_fixture):
     tmp_path = import_fixture