1use crate::PythonAnalysisError;
2use crate::parser::parse_python;
3use tree_sitter::Tree;
4
5pub 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
21pub 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 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 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}