codegraph_python/entities/
trait_.rs

1use serde::{Deserialize, Serialize};
2
3/// Represents a Python Protocol or Abstract Base Class (trait-like entity)
4#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
5pub struct TraitEntity {
6    /// Trait/Protocol name
7    pub name: String,
8
9    /// Visibility: "public" or "private"
10    pub visibility: String,
11
12    /// Starting line number (1-indexed)
13    pub line_start: usize,
14
15    /// Ending line number (1-indexed)
16    pub line_end: usize,
17
18    /// Is this a Protocol (typing.Protocol)?
19    pub is_protocol: bool,
20
21    /// Is this an Abstract Base Class (abc.ABC)?
22    pub is_abc: bool,
23
24    /// Method signatures required by this trait
25    pub required_methods: Vec<String>,
26
27    /// Documentation string (docstring)
28    pub doc_comment: Option<String>,
29
30    /// Decorators applied
31    pub attributes: Vec<String>,
32}
33
34impl TraitEntity {
35    /// Create a new trait entity
36    pub fn new(name: impl Into<String>, line_start: usize, line_end: usize) -> Self {
37        let name = name.into();
38        let visibility = if name.starts_with('_') {
39            "private".to_string()
40        } else {
41            "public".to_string()
42        };
43
44        Self {
45            name,
46            visibility,
47            line_start,
48            line_end,
49            is_protocol: false,
50            is_abc: false,
51            required_methods: Vec::new(),
52            doc_comment: None,
53            attributes: Vec::new(),
54        }
55    }
56
57    /// Set as a Protocol
58    pub fn set_protocol(mut self, is_protocol: bool) -> Self {
59        self.is_protocol = is_protocol;
60        self
61    }
62
63    /// Set as an ABC
64    pub fn set_abc(mut self, is_abc: bool) -> Self {
65        self.is_abc = is_abc;
66        self
67    }
68
69    /// Add a required method
70    pub fn add_required_method(mut self, method: impl Into<String>) -> Self {
71        self.required_methods.push(method.into());
72        self
73    }
74
75    /// Set the docstring
76    pub fn set_doc_comment(mut self, doc: Option<String>) -> Self {
77        self.doc_comment = doc;
78        self
79    }
80
81    /// Add a decorator
82    pub fn add_attribute(mut self, attr: impl Into<String>) -> Self {
83        self.attributes.push(attr.into());
84        self
85    }
86}
87
88#[cfg(test)]
89mod tests {
90    use super::*;
91
92    #[test]
93    fn test_trait_entity_new() {
94        let trait_entity = TraitEntity::new("Serializable", 1, 10);
95        assert_eq!(trait_entity.name, "Serializable");
96        assert_eq!(trait_entity.visibility, "public");
97        assert!(!trait_entity.is_protocol);
98        assert!(!trait_entity.is_abc);
99    }
100
101    #[test]
102    fn test_trait_entity_protocol() {
103        let protocol = TraitEntity::new("Drawable", 1, 10)
104            .set_protocol(true)
105            .add_required_method("draw");
106
107        assert!(protocol.is_protocol);
108        assert_eq!(protocol.required_methods, vec!["draw"]);
109    }
110
111    #[test]
112    fn test_trait_entity_abc() {
113        let abc = TraitEntity::new("BaseClass", 1, 10)
114            .set_abc(true)
115            .add_required_method("abstract_method");
116
117        assert!(abc.is_abc);
118        assert_eq!(abc.required_methods, vec!["abstract_method"]);
119    }
120}
codegraph_python/entities/trait_.rs

codegraph_python/entities/
trait_.rs
