use crate::config::CodeAgentConfig;
use crate::error::CodeAgentError;
use crate::parser::CodeParser;
use std::collections::HashMap;
use tree_sitter::{Node, Tree};
pub struct CodeDocumenter {
parser: CodeParser,
config: CodeAgentConfig,
docs: HashMap<String, String>,
}
impl CodeDocumenter {
pub fn new(config: CodeAgentConfig) -> Result<Self, CodeAgentError> {
let parser = CodeParser::new(config.clone())?;
Ok(Self {
parser,
config,
docs: HashMap::new(),
})
}
pub fn document_file(&mut self, path: &str) -> Result<(), CodeAgentError> {
self.parser.parse_file(path.as_ref())?;
let tree = self
.parser
.parser_mut()
.parse(path, None)
.ok_or_else(|| CodeAgentError::Parser("Failed to parse file".to_string()))?;
self.document_tree(&tree)
}
pub fn document_content(&mut self, content: &str) -> Result<(), CodeAgentError> {
self.parser.parse_content(content)?;
let tree = self
.parser
.parser_mut()
.parse(content.as_bytes(), None)
.ok_or_else(|| CodeAgentError::Parser("Failed to parse content".to_string()))?;
self.document_tree(&tree)
}
fn document_tree(&mut self, tree: &Tree) -> Result<(), CodeAgentError> {
let root_node = tree.root_node();
if self.config.enable_documentation_generation {
self.generate_documentation(&root_node)?;
}
if self.config.enable_test_generation {
self.generate_tests(&root_node)?;
}
Ok(())
}
fn generate_documentation(&mut self, node: &Node) -> Result<(), CodeAgentError> {
match node.kind() {
"function_item" => self.document_function(node)?,
"struct_item" => self.document_struct(node)?,
"enum_item" => self.document_enum(node)?,
"trait_item" => self.document_trait(node)?,
"impl_item" => self.document_impl(node)?,
"module_item" => self.document_module(node)?,
_ => {}
}
for i in 0..node.child_count() {
if let Some(child) = node.child(i) {
self.generate_documentation(&child)?;
}
}
Ok(())
}
fn document_function(&mut self, node: &Node) -> Result<(), CodeAgentError> {
Ok(())
}
fn document_struct(&mut self, node: &Node) -> Result<(), CodeAgentError> {
Ok(())
}
fn document_enum(&mut self, node: &Node) -> Result<(), CodeAgentError> {
Ok(())
}
fn document_trait(&mut self, node: &Node) -> Result<(), CodeAgentError> {
Ok(())
}
fn document_impl(&mut self, node: &Node) -> Result<(), CodeAgentError> {
Ok(())
}
fn document_module(&mut self, node: &Node) -> Result<(), CodeAgentError> {
Ok(())
}
fn generate_tests(&mut self, node: &Node) -> Result<(), CodeAgentError> {
match node.kind() {
"function_item" => self.generate_function_tests(node)?,
"struct_item" => self.generate_struct_tests(node)?,
"enum_item" => self.generate_enum_tests(node)?,
"trait_item" => self.generate_trait_tests(node)?,
"impl_item" => self.generate_impl_tests(node)?,
_ => {}
}
for i in 0..node.child_count() {
if let Some(child) = node.child(i) {
self.generate_tests(&child)?;
}
}
Ok(())
}
fn generate_function_tests(&mut self, node: &Node) -> Result<(), CodeAgentError> {
Ok(())
}
fn generate_struct_tests(&mut self, node: &Node) -> Result<(), CodeAgentError> {
Ok(())
}
fn generate_enum_tests(&mut self, node: &Node) -> Result<(), CodeAgentError> {
Ok(())
}
fn generate_trait_tests(&mut self, node: &Node) -> Result<(), CodeAgentError> {
Ok(())
}
fn generate_impl_tests(&mut self, node: &Node) -> Result<(), CodeAgentError> {
Ok(())
}
pub fn docs(&self) -> &HashMap<String, String> {
&self.docs
}
pub fn docs_mut(&mut self) -> &mut HashMap<String, String> {
&mut self.docs
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_documenter_creation() {
let config = CodeAgentConfig::default();
let documenter = CodeDocumenter::new(config);
assert!(documenter.is_ok());
}
#[test]
fn test_document_content() {
let config = CodeAgentConfig::default();
let mut documenter = CodeDocumenter::new(config).unwrap();
let result = documenter.document_content("fn main() { println!(\"Hello, world!\"); }");
assert!(result.is_ok());
}
}