Skip to main content

code_analyze_core/languages/
python.rs

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