From 5dc106f5721c8382870744bf1f1a8b8d65e5f3e4 Mon Sep 17 00:00:00 2001
From: Alex Povel <git@alexpovel.de>
Date: Tue, 30 Jul 2024 19:19:16 +0200
Subject: [PATCH] feat(python): Scope type hints

---
 README.md                                     |  1 +
 src/scoping/langs/python.rs                   |  3 +++
 tests/langs/mod.rs                            |  5 +++++
 .../r#mod__langs__base.py_types.snap          | 19 +++++++++++++++++++
 4 files changed, 28 insertions(+)
 create mode 100644 tests/langs/snapshots/r#mod__langs__base.py_types.snap

diff --git a/README.md b/README.md
index 967adada..cf915bac 100644
--- a/README.md
+++ b/README.md
@@ -1434,6 +1434,7 @@ Language scopes:
           - globals:              Global, i.e. module-level variables
           - variable-identifiers: Identifiers for variables (left-hand side of
             assignments)
+          - types:                Types in type hints
 
       --python-query <PYTHON_QUERY>
           Scope Python code using a custom tree-sitter query.
diff --git a/src/scoping/langs/python.rs b/src/scoping/langs/python.rs
index 6898e020..5ce83544 100644
--- a/src/scoping/langs/python.rs
+++ b/src/scoping/langs/python.rs
@@ -47,6 +47,8 @@ pub enum PreparedPythonQuery {
     Globals,
     /// Identifiers for variables (left-hand side of assignments).
     VariableIdentifiers,
+    /// Types in type hints.
+    Types,
 }
 
 impl From<PreparedPythonQuery> for TSQuery {
@@ -162,6 +164,7 @@ impl From<PreparedPythonQuery> for TSQuery {
                 PreparedPythonQuery::VariableIdentifiers => {
                     "(assignment left: (identifier) @identifier)"
                 }
+                PreparedPythonQuery::Types => "(type) @type",
             },
         )
         .expect("Prepared queries to be valid")
diff --git a/tests/langs/mod.rs b/tests/langs/mod.rs
index 17dc35d5..f64c5cc5 100644
--- a/tests/langs/mod.rs
+++ b/tests/langs/mod.rs
@@ -157,6 +157,11 @@ impl InScopeLinePart {
     include_str!("python/base.py"),
     Python::new(CodeQuery::Prepared(PreparedPythonQuery::VariableIdentifiers)),
 )]
+#[case(
+    "base.py_types",
+    include_str!("python/base.py"),
+    Python::new(CodeQuery::Prepared(PreparedPythonQuery::Types)),
+)]
 #[case(
     "base.ts_strings",
     include_str!("typescript/base.ts"),
diff --git a/tests/langs/snapshots/r#mod__langs__base.py_types.snap b/tests/langs/snapshots/r#mod__langs__base.py_types.snap
new file mode 100644
index 00000000..a783b47d
--- /dev/null
+++ b/tests/langs/snapshots/r#mod__langs__base.py_types.snap
@@ -0,0 +1,19 @@
+---
+source: tests/langs/mod.rs
+expression: inscope_parts
+---
+- n: 16
+  l: "test_var: int = 10\n"
+  m: "          ^^^       "
+- n: 63
+  l: "    def class_method(cls) -> None:\n"
+  m: "                             ^^^^   "
+- n: 69
+  l: "    def instance_method(self) -> None:\n"
+  m: "                                 ^^^^   "
+- n: 76
+  l: "    def static_method() -> None:\n"
+  m: "                           ^^^^   "
+- n: 112
+  l: "def exception_handling(x) -> None:\n"
+  m: "                             ^^^^   "