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, Parser};
6
7use crate::core::{Kind, Symbol};
8use crate::lang::LanguagePlugin;
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        let mut parser = Parser::new();
25        if parser
26            .set_language(&tree_sitter_python::LANGUAGE.into())
27            .is_err()
28        {
29            return Vec::new();
30        }
31        let Some(tree) = parser.parse(source, None) else {
32            return Vec::new();
33        };
34        let mut out = Vec::new();
35        let ctx = Ctx {
36            src: source.as_bytes(),
37            file,
38        };
39        ctx.walk(tree.root_node(), None, false, &mut out);
40        out
41    }
42}
43
44struct Ctx<'a> {
45    src: &'a [u8],
46    file: &'a str,
47}
48
49impl Ctx<'_> {
50    /// `parent` is the enclosing qualified name; `in_class` is true inside a
51    /// class body, where a `def` is a method.
52    fn walk(&self, node: Node, parent: Option<&str>, in_class: bool, out: &mut Vec<Symbol>) {
53        let mut cursor = node.walk();
54        for child in node.children(&mut cursor) {
55            match child.kind() {
56                "class_definition" => {
57                    if let Some(name) = self.field_text(child, "name") {
58                        out.push(self.symbol(&name, Kind::Class, child, parent));
59                        let qualified = qualify(parent, &name);
60                        self.walk(child, Some(&qualified), true, out);
61                    }
62                }
63                "function_definition" => {
64                    if let Some(name) = self.field_text(child, "name") {
65                        let kind = if in_class {
66                            Kind::Method
67                        } else {
68                            Kind::Function
69                        };
70                        out.push(self.symbol(&name, kind, child, parent));
71                    }
72                    // don't descend into a def body (nested defs rarely navigated)
73                }
74                // a decorated class/function: descend so the wrapped def is seen
75                // in the same context
76                _ => self.walk(child, parent, in_class, out),
77            }
78        }
79    }
80
81    fn field_text(&self, node: Node, field: &str) -> Option<String> {
82        node.child_by_field_name(field)
83            .and_then(|n| n.utf8_text(self.src).ok())
84            .map(str::to_string)
85    }
86
87    fn symbol(&self, name: &str, kind: Kind, node: Node, parent: Option<&str>) -> Symbol {
88        Symbol {
89            name: name.to_string(),
90            kind,
91            language: LANGUAGE.to_string(),
92            file: self.file.to_string(),
93            line: node.start_position().row as u32 + 1,
94            parent: parent.map(str::to_string),
95        }
96    }
97}
98
99fn qualify(parent: Option<&str>, name: &str) -> String {
100    match parent {
101        Some(p) => format!("{p}.{name}"),
102        None => name.to_string(),
103    }
104}
105
106#[cfg(test)]
107mod tests {
108    use super::*;
109
110    fn extract(source: &str) -> Vec<Symbol> {
111        Python.extract("test.py", source)
112    }
113
114    fn find<'a>(syms: &'a [Symbol], name: &str) -> &'a Symbol {
115        syms.iter()
116            .find(|s| s.name == name)
117            .unwrap_or_else(|| panic!("no symbol named {name} in {syms:?}"))
118    }
119
120    #[test]
121    fn extracts_classes_methods_and_functions() {
122        let src = r#"
123class Account:
124    def deposit(self, amount):
125        pass
126
127    @property
128    def balance(self):
129        return 0
130
131def build():
132    return Account()
133"#;
134        let syms = extract(src);
135
136        let account = find(&syms, "Account");
137        assert_eq!(account.kind, Kind::Class);
138        assert_eq!(account.parent, None);
139
140        let deposit = find(&syms, "deposit");
141        assert_eq!(deposit.kind, Kind::Method);
142        assert_eq!(deposit.parent.as_deref(), Some("Account"));
143
144        // a decorated method is still found, still a method
145        assert_eq!(find(&syms, "balance").kind, Kind::Method);
146
147        // a module-level def is a function
148        let build = find(&syms, "build");
149        assert_eq!(build.kind, Kind::Function);
150        assert_eq!(build.parent, None);
151
152        assert_eq!(account.language, "python");
153    }
154
155    #[test]
156    fn empty_and_unparseable_yield_no_symbols() {
157        assert!(extract("").is_empty());
158        assert!(extract("# just a comment\n").is_empty());
159    }
160}
reference_query/lang/python/mod.rs

reference_query/lang/python/
mod.rs
