Skip to main content

code_analyze_mcp/languages/
python.rs

1/// Tree-sitter query for extracting Python elements (functions and classes).
2pub const ELEMENT_QUERY: &str = r"
3(function_definition
4  name: (identifier) @func_name) @function
5(class_definition
6  name: (identifier) @class_name) @class
7";
8
9/// Tree-sitter query for extracting function calls.
10pub const CALL_QUERY: &str = r"
11(call
12  function: (identifier) @call)
13(call
14  function: (attribute attribute: (identifier) @call))
15";
16
17/// Tree-sitter query for extracting type references.
18/// Python grammar has no `type_identifier` node; use `(type (identifier) @type_ref)`
19/// to capture type names in annotations and `generic_type` for parameterized types.
20pub const REFERENCE_QUERY: &str = r"
21(type (identifier) @type_ref)
22(generic_type (identifier) @type_ref)
23";
24
25/// Tree-sitter query for extracting Python imports.
26pub const IMPORT_QUERY: &str = r"
27(import_statement) @import_path
28(import_from_statement) @import_path
29";
30
31use tree_sitter::Node;
32
33/// Extract inheritance information from a Python class node.
34#[must_use]
35pub fn extract_inheritance(node: &Node, source: &str) -> Vec<String> {
36    let mut inherits = Vec::new();
37
38    // Get superclasses field from class_definition
39    if let Some(superclasses) = node.child_by_field_name("superclasses") {
40        // superclasses contains an argument_list
41        for i in 0..superclasses.named_child_count() {
42            if let Some(child) = superclasses.named_child(u32::try_from(i).unwrap_or(u32::MAX))
43                && matches!(child.kind(), "identifier" | "attribute")
44            {
45                let text = &source[child.start_byte()..child.end_byte()];
46                inherits.push(text.to_string());
47            }
48        }
49    }
50
51    inherits
52}