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.
34pub fn extract_inheritance(node: &Node, source: &str) -> Vec<String> {
35    let mut inherits = Vec::new();
36
37    // Get superclasses field from class_definition
38    if let Some(superclasses) = node.child_by_field_name("superclasses") {
39        // superclasses contains an argument_list
40        for i in 0..superclasses.named_child_count() {
41            if let Some(child) = superclasses.named_child(i as u32)
42                && matches!(child.kind(), "identifier" | "attribute")
43            {
44                let text = &source[child.start_byte()..child.end_byte()];
45                inherits.push(text.to_string());
46            }
47        }
48    }
49
50    inherits
51}