codebank/parser/formatter/
rules.rs

1use crate::parser::LanguageType;
2
3#[derive(Debug, Clone, Copy)]
4#[allow(dead_code)]
5pub struct FormatterRules {
6    pub summary_ellipsis: &'static str,
7    pub field_sep: &'static str,
8    pub function_body_start_marker: &'static str,
9    pub function_body_end_marker: &'static str,
10    pub doc_marker: &'static str,
11    pub test_markers: &'static [&'static str],
12    pub test_module_markers: &'static [&'static str],
13}
14
15const RUST_RULES: FormatterRules = FormatterRules {
16    summary_ellipsis: " { ... }",
17    field_sep: ",",
18    function_body_start_marker: "{",
19    function_body_end_marker: "}",
20    doc_marker: "///",
21    test_markers: &["#[test]", "#[cfg(test)]"],
22    test_module_markers: &["#[cfg(test)]", "tests"],
23};
24
25const PYTHON_RULES: FormatterRules = FormatterRules {
26    summary_ellipsis: ": ...",
27    field_sep: "",
28    function_body_start_marker: ":",
29    function_body_end_marker: "",
30    doc_marker: "#",
31    test_markers: &["@pytest", "test_"],
32    test_module_markers: &["test_"],
33};
34
35const TS_RULES: FormatterRules = FormatterRules {
36    summary_ellipsis: " { ... }",
37    field_sep: ",",
38    function_body_start_marker: "{",
39    function_body_end_marker: "}",
40    doc_marker: "//",
41    test_markers: &["@test", "test_"],
42    test_module_markers: &["test_"],
43};
44
45const C_RULES: FormatterRules = FormatterRules {
46    summary_ellipsis: " { ... }",
47    field_sep: ",",
48    function_body_start_marker: "{",
49    function_body_end_marker: "}",
50    doc_marker: "//",
51    test_markers: &["@test", "test_"],
52    test_module_markers: &["test_"],
53};
54
55const GO_RULES: FormatterRules = FormatterRules {
56    summary_ellipsis: " { ... }",
57    field_sep: ",",
58    function_body_start_marker: "{",
59    function_body_end_marker: "}",
60    doc_marker: "//",
61    test_markers: &["test_"],
62    test_module_markers: &["test_"],
63};
64
65const UNKNOWN_RULES: FormatterRules = FormatterRules {
66    summary_ellipsis: "...",
67    field_sep: "",
68    function_body_start_marker: "",
69    function_body_end_marker: "",
70    doc_marker: "//",
71    test_markers: &[],
72    test_module_markers: &[],
73};
74
75impl FormatterRules {
76    #[inline(always)]
77    pub fn for_language(lang: LanguageType) -> Self {
78        match lang {
79            LanguageType::Rust => RUST_RULES,
80            LanguageType::Python => PYTHON_RULES,
81            LanguageType::TypeScript => TS_RULES,
82            LanguageType::Cpp => C_RULES,
83            LanguageType::Go => GO_RULES,
84            LanguageType::Unknown => UNKNOWN_RULES,
85        }
86    }
87
88    pub fn is_test_function(&self, attributes: &[String]) -> bool {
89        attributes
90            .iter()
91            .any(|attr| self.test_markers.iter().any(|marker| attr.contains(marker)))
92    }
93
94    pub fn is_test_module(&self, name: &str, attributes: &[String]) -> bool {
95        self.test_module_markers.iter().any(|marker| {
96            name.starts_with(marker) || attributes.iter().any(|attr| attr.contains(marker))
97        })
98    }
99
100    pub fn format_signature(&self, source: &str, signature: Option<&str>) -> String {
101        let sig_to_use = signature.unwrap_or(source).trim();
102
103        // Find the start of the body marker if it exists
104        let body_start_marker_pos = if !self.function_body_start_marker.is_empty() {
105            sig_to_use.find(self.function_body_start_marker)
106        } else {
107            None
108        };
109
110        let clean_sig = if let Some(idx) = body_start_marker_pos {
111            sig_to_use[0..idx].trim_end()
112        } else {
113            sig_to_use
114        };
115
116        // Append the language-specific summary ellipsis
117        format!("{}{}", clean_sig, self.summary_ellipsis)
118    }
119}
120
121#[cfg(test)]
122mod tests {
123    use super::*;
124
125    #[test]
126    fn test_rust_rules() {
127        let rules = FormatterRules::for_language(LanguageType::Rust);
128        assert_eq!(rules.summary_ellipsis, " { ... }");
129        assert_eq!(rules.function_body_start_marker, "{");
130        assert_eq!(rules.test_markers, &["#[test]", "#[cfg(test)]"]);
131        assert_eq!(rules.test_module_markers, &["#[cfg(test)]", "tests"]);
132    }
133
134    #[test]
135    fn test_python_rules() {
136        let rules = FormatterRules::for_language(LanguageType::Python);
137        assert_eq!(rules.summary_ellipsis, ": ...");
138        assert_eq!(rules.function_body_start_marker, ":");
139        assert_eq!(rules.test_markers, &["@pytest", "test_"]);
140        assert_eq!(rules.test_module_markers, &["test_"]);
141    }
142
143    #[test]
144    fn test_unknown_language_rules() {
145        let rules = FormatterRules::for_language(LanguageType::Unknown);
146        assert_eq!(rules.summary_ellipsis, "...");
147        assert_eq!(rules.function_body_start_marker, "");
148        assert!(rules.test_markers.is_empty());
149        assert!(rules.test_module_markers.is_empty());
150    }
151
152    #[test]
153    fn test_is_test_function() {
154        let rules = FormatterRules::for_language(LanguageType::Rust);
155
156        // Test Rust test function detection
157        assert!(rules.is_test_function(&["#[test]".to_string()]));
158        assert!(rules.is_test_function(&["#[cfg(test)]".to_string()]));
159        assert!(!rules.is_test_function(&["#[derive(Debug)]".to_string()]));
160
161        let rules = FormatterRules::for_language(LanguageType::Python);
162
163        // Test Python test function detection
164        assert!(rules.is_test_function(&["@pytest.mark.test".to_string()]));
165        assert!(rules.is_test_function(&["test_function".to_string()]));
166        assert!(!rules.is_test_function(&["regular_function".to_string()]));
167    }
168
169    #[test]
170    fn test_is_test_module() {
171        let rules = FormatterRules::for_language(LanguageType::Rust);
172
173        // Test Rust test module detection
174        assert!(rules.is_test_module("tests", &[]));
175        assert!(rules.is_test_module("module", &["#[cfg(test)]".to_string()]));
176        assert!(!rules.is_test_module("module", &[]));
177
178        let rules = FormatterRules::for_language(LanguageType::Python);
179
180        // Test Python test module detection
181        assert!(rules.is_test_module("test_module", &[]));
182        assert!(!rules.is_test_module("regular_module", &[]));
183    }
184
185    #[test]
186    fn test_format_signature() {
187        let rules = FormatterRules::for_language(LanguageType::Rust);
188
189        // Test with signature provided
190        assert_eq!(
191            rules.format_signature("fn test() {}", Some("fn test()")),
192            "fn test() { ... }"
193        );
194
195        // Test without signature, with body start marker
196        assert_eq!(
197            rules.format_signature("fn test() {", None),
198            "fn test() { ... }"
199        );
200
201        // Test without signature or body marker (e.g., trait method)
202        assert_eq!(
203            rules.format_signature("fn test()", None),
204            "fn test() { ... }"
205        );
206
207        // Test with extra whitespace
208        assert_eq!(
209            rules.format_signature("fn test()  {", None),
210            "fn test() { ... }"
211        );
212
213        let rules = FormatterRules::for_language(LanguageType::Python);
214
215        // Test Python function signature
216        assert_eq!(
217            rules.format_signature("def test():", None),
218            "def test(): ..."
219        );
220        assert_eq!(
221            rules.format_signature("def test()", None), // No colon
222            "def test(): ..."
223        );
224    }
225}
codebank/parser/formatter/rules.rs

codebank/parser/formatter/
rules.rs
