probe_code/language/
python.rs

1use super::language_trait::LanguageImpl;
2use tree_sitter::{Language as TSLanguage, Node};
3
4/// Implementation of LanguageImpl for Python
5pub struct PythonLanguage;
6
7impl Default for PythonLanguage {
8    fn default() -> Self {
9        Self::new()
10    }
11}
12
13impl PythonLanguage {
14    pub fn new() -> Self {
15        PythonLanguage
16    }
17}
18
19impl LanguageImpl for PythonLanguage {
20    fn get_tree_sitter_language(&self) -> TSLanguage {
21        tree_sitter_python::LANGUAGE.into()
22    }
23
24    fn get_extension(&self) -> &'static str {
25        "py"
26    }
27
28    fn is_acceptable_parent(&self, node: &Node) -> bool {
29        matches!(node.kind(), "function_definition" | "class_definition")
30    }
31
32    fn is_test_node(&self, node: &Node, source: &[u8]) -> bool {
33        let debug_mode = std::env::var("DEBUG").unwrap_or_default() == "1";
34        let node_type = node.kind();
35
36        // Python: Check function_definition nodes with names starting with test_
37        if node_type == "function_definition" {
38            let mut cursor = node.walk();
39            for child in node.children(&mut cursor) {
40                if child.kind() == "identifier" {
41                    let name = child.utf8_text(source).unwrap_or("");
42                    if name.starts_with("test_") {
43                        if debug_mode {
44                            println!("DEBUG: Test node detected (Python): test_ function");
45                        }
46                        return true;
47                    }
48                }
49            }
50        }
51
52        false
53    }
54}