reference_query/lang/python/
mod.rs

1//! Python plugin. Extracts `class` → class and `def` (free → function, inside a
2//! class → method), qualified with `.` (`method · Account`, nested class
3//! `Inner · Outer`). Decorators are transparent — the wrapped def is what counts.
4
5use tree_sitter::Node;
6
7use crate::core::{Kind, Symbol};
8use crate::lang::{Ctx, LanguagePlugin, extract_with, qualify};
9
10const LANGUAGE: &str = "python";
11
12pub struct Python;
13
14impl LanguagePlugin for Python {
15    fn language(&self) -> &'static str {
16        LANGUAGE
17    }
18
19    fn extensions(&self) -> &[&str] {
20        &["py"]
21    }
22
23    fn extract(&self, file: &str, source: &str) -> Vec<Symbol> {
24        extract_with(
25            LANGUAGE,
26            tree_sitter_python::LANGUAGE.into(),
27            file,
28            source,
29            |ctx, root, out| walk(ctx, root, None, false, out),
30        )
31    }
32}
33
34/// `parent` is the enclosing qualified name; `in_class` is true inside a
35/// class body, where a `def` is a method.
36fn walk(ctx: &Ctx, node: Node, parent: Option<&str>, in_class: bool, out: &mut Vec<Symbol>) {
37    let mut cursor = node.walk();
38    for child in node.children(&mut cursor) {
39        match child.kind() {
40            "class_definition" => {
41                if let Some(name) = ctx.field_text(child, "name") {
42                    let mut s = ctx.symbol(&name, Kind::Class, child, parent);
43                    s.visibility = Some(name_visibility(&name));
44                    out.push(s);
45                    let qualified = qualify(parent, &name, ".");
46                    walk(ctx, child, Some(&qualified), true, out);
47                }
48            }
49            "function_definition" => {
50                if let Some(name) = ctx.field_text(child, "name") {
51                    let kind = if in_class {
52                        Kind::Method
53                    } else {
54                        Kind::Function
55                    };
56                    let mut s = ctx.symbol(&name, kind, child, parent);
57                    s.visibility = Some(name_visibility(&name));
58                    out.push(s);
59                }
60                // don't descend into a def body (nested defs rarely navigated)
61            }
62            // a decorated class/function: descend so the wrapped def is seen
63            // in the same context
64            _ => walk(ctx, child, parent, in_class, out),
65        }
66    }
67}
68
69/// Python's naming convention: a leading underscore marks internal —
70/// except dunders (`__init__`), which are the public protocol surface.
71fn name_visibility(name: &str) -> &'static str {
72    let dunder = name.starts_with("__") && name.ends_with("__");
73    if name.starts_with('_') && !dunder {
74        "private"
75    } else {
76        "public"
77    }
78}
79
80#[cfg(test)]
81mod tests {
82    use super::*;
83
84    fn extract(source: &str) -> Vec<Symbol> {
85        Python.extract("test.py", source)
86    }
87
88    fn find<'a>(syms: &'a [Symbol], name: &str) -> &'a Symbol {
89        syms.iter()
90            .find(|s| s.name == name)
91            .unwrap_or_else(|| panic!("no symbol named {name} in {syms:?}"))
92    }
93
94    #[test]
95    fn extracts_classes_methods_and_functions() {
96        let src = r#"
97class Account:
98    def deposit(self, amount):
99        pass
100
101    @property
102    def balance(self):
103        return 0
104
105def build():
106    return Account()
107"#;
108        let syms = extract(src);
109
110        let account = find(&syms, "Account");
111        assert_eq!(account.kind, Kind::Class);
112        assert_eq!(account.parent, None);
113
114        let deposit = find(&syms, "deposit");
115        assert_eq!(deposit.kind, Kind::Method);
116        assert_eq!(deposit.parent.as_deref(), Some("Account"));
117
118        // a decorated method is still found, still a method
119        assert_eq!(find(&syms, "balance").kind, Kind::Method);
120
121        // a module-level def is a function
122        let build = find(&syms, "build");
123        assert_eq!(build.kind, Kind::Function);
124        assert_eq!(build.parent, None);
125
126        assert_eq!(account.language, "python");
127    }
128
129    #[test]
130    fn empty_and_unparseable_yield_no_symbols() {
131        assert!(extract("").is_empty());
132        assert!(extract("# just a comment\n").is_empty());
133    }
134
135    #[test]
136    fn underscore_names_read_as_private_except_dunders() {
137        let src = "class Account:\n    def _internal(self):\n        pass\n    def __init__(self):\n        pass\n\ndef fetch():\n    pass\n";
138        let syms = extract(src);
139        assert_eq!(find(&syms, "_internal").visibility, Some("private"));
140        assert_eq!(find(&syms, "__init__").visibility, Some("public"));
141        assert_eq!(find(&syms, "fetch").visibility, Some("public"));
142    }
143}
reference_query/lang/python/mod.rs

reference_query/lang/python/
mod.rs
