Skip to main content

pecto_python/
context.rs

1use crate::PythonAnalysisError;
2use crate::parser::parse_python;
3use tree_sitter::Tree;
4
5/// A parsed Python source file ready for extraction.
6pub struct ParsedFile {
7    pub path: String,
8    pub source: String,
9    pub tree: Tree,
10}
11
12impl ParsedFile {
13    pub fn parse(source: String, path: String) -> Result<Self, PythonAnalysisError> {
14        let tree = parse_python(&source).ok_or_else(|| {
15            PythonAnalysisError::ParseError(path.clone(), "Failed to parse Python".to_string())
16        })?;
17        Ok(Self { path, source, tree })
18    }
19}
20
21/// Holds all parsed files for cross-file lookups during extraction.
22pub struct AnalysisContext {
23    pub files: Vec<ParsedFile>,
24}
25
26impl AnalysisContext {
27    pub fn new(files: Vec<ParsedFile>) -> Self {
28        Self { files }
29    }
30
31    /// Find a class definition by name across all parsed files.
32    pub fn find_class_by_name(&self, class_name: &str) -> Option<&ParsedFile> {
33        use crate::extractors::common::node_text;
34
35        for file in &self.files {
36            let root = file.tree.root_node();
37            let source = file.source.as_bytes();
38            for i in 0..root.named_child_count() {
39                let node = root.named_child(i).unwrap();
40                // Python classes can be top-level or inside decorated_definition
41                let class_node = if node.kind() == "class_definition" {
42                    Some(node)
43                } else if node.kind() == "decorated_definition" {
44                    node.named_children(&mut node.walk())
45                        .find(|c| c.kind() == "class_definition")
46                } else {
47                    None
48                };
49
50                if let Some(cn) = class_node
51                    && let Some(name) = cn.child_by_field_name("name")
52                    && node_text(&name, source) == class_name
53                {
54                    return Some(file);
55                }
56            }
57        }
58        None
59    }
60}