codegraph_python/entities/
file.rs

1use serde::{Deserialize, Serialize};
2
3/// Represents a Python module/file
4#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
5pub struct ModuleEntity {
6    /// Module name (derived from file path)
7    pub name: String,
8
9    /// Visibility (always "public" for modules)
10    pub visibility: String,
11
12    /// Starting line number (always 1 for modules)
13    pub line_start: usize,
14
15    /// Ending line number (last line of file)
16    pub line_end: usize,
17
18    /// Module docstring
19    pub doc_comment: Option<String>,
20}
21
22impl ModuleEntity {
23    /// Create a new module entity
24    pub fn new(name: impl Into<String>, line_end: usize) -> Self {
25        Self {
26            name: name.into(),
27            visibility: "public".to_string(),
28            line_start: 1,
29            line_end,
30            doc_comment: None,
31        }
32    }
33
34    /// Set the module docstring
35    pub fn set_doc_comment(mut self, doc: Option<String>) -> Self {
36        self.doc_comment = doc;
37        self
38    }
39
40    /// Get the total number of lines in the module
41    pub fn line_count(&self) -> usize {
42        self.line_end
43    }
44}
45
46#[cfg(test)]
47mod tests {
48    use super::*;
49
50    #[test]
51    fn test_module_entity_new() {
52        let module = ModuleEntity::new("my_module", 100);
53        assert_eq!(module.name, "my_module");
54        assert_eq!(module.visibility, "public");
55        assert_eq!(module.line_start, 1);
56        assert_eq!(module.line_end, 100);
57        assert_eq!(module.line_count(), 100);
58    }
59
60    #[test]
61    fn test_module_entity_with_docstring() {
62        let module = ModuleEntity::new("my_module", 50)
63            .set_doc_comment(Some("This is a module".to_string()));
64
65        assert_eq!(module.doc_comment, Some("This is a module".to_string()));
66    }
67}