jsonschema_annotator/annotator/
toml.rs

1use toml_edit::{DocumentMut, Item, Table};
2
3use super::{Annotator, AnnotatorConfig, ExistingCommentBehavior};
4use crate::error::{AnnotatorError, AnnotatorErrorKind, Error};
5use crate::schema::{Annotation, AnnotationMap};
6
7/// TOML document annotator using toml_edit
8pub struct TomlAnnotator {
9    config: AnnotatorConfig,
10}
11
12impl TomlAnnotator {
13    pub fn new(config: AnnotatorConfig) -> Self {
14        Self { config }
15    }
16
17    fn format_comment(&self, annotation: &Annotation) -> Option<String> {
18        let mut lines = Vec::new();
19
20        if self.config.include_title {
21            if let Some(title) = &annotation.title {
22                lines.push(format!("# {}", title));
23            }
24        }
25
26        if self.config.include_description {
27            if let Some(desc) = &annotation.description {
28                let width = self.config.max_line_width.unwrap_or(78);
29                for line in textwrap::wrap(desc, width.saturating_sub(2)) {
30                    lines.push(format!("# {}", line));
31                }
32            }
33        }
34
35        if self.config.include_default {
36            if let Some(default) = &annotation.default {
37                lines.push(format!("# Default: {}", default));
38            }
39        }
40
41        if lines.is_empty() {
42            None
43        } else {
44            // Add newline after comments so it appears before the key
45            Some(lines.join("\n") + "\n")
46        }
47    }
48
49    fn annotate_table(
50        &self,
51        table: &mut Table,
52        path: &[String],
53        annotations: &AnnotationMap,
54    ) {
55        // Collect keys first to avoid borrow issues
56        // Use deref to str to get the key string (Key derefs to str)
57        let keys: Vec<String> = table.iter().map(|(k, _)| (*k).to_string()).collect();
58
59        for key_string in keys {
60            let mut current_path = path.to_vec();
61            current_path.push(key_string.clone());
62            let path_string = current_path.join(".");
63
64            // Get mutable access to the key-value pair
65            if let Some((mut key, item)) = table.get_key_value_mut(&key_string) {
66                // Handle tables vs regular values differently
67                match item {
68                    Item::Table(nested) => {
69                        // For tables, use the table's own decor (appears before the [header])
70                        if let Some(ann) = annotations.get(&path_string) {
71                            if let Some(comment) = self.format_comment(ann) {
72                                let decor = nested.decor_mut();
73                                let existing = decor.prefix().map(|s| s.as_str().unwrap_or("")).unwrap_or("");
74                                let has_existing = existing.trim().starts_with('#');
75
76                                let new_prefix = match self.config.existing_comments {
77                                    ExistingCommentBehavior::Skip if has_existing => None,
78                                    ExistingCommentBehavior::Prepend if has_existing => {
79                                        Some(format!("{}{}", comment, existing))
80                                    }
81                                    ExistingCommentBehavior::Append if has_existing => {
82                                        Some(format!("{}{}", existing, comment))
83                                    }
84                                    _ => Some(comment), // Replace or no existing comment
85                                };
86
87                                if let Some(prefix) = new_prefix {
88                                    decor.set_prefix(prefix);
89                                }
90                            }
91                        }
92                        // Recurse into nested tables
93                        self.annotate_table(nested, &current_path, annotations);
94                    }
95                    Item::Value(toml_edit::Value::InlineTable(_)) => {
96                        // Can't easily modify inline tables, skip for now
97                    }
98                    _ => {
99                        // For regular values, use the key's decor
100                        if let Some(ann) = annotations.get(&path_string) {
101                            if let Some(comment) = self.format_comment(ann) {
102                                let decor = key.leaf_decor_mut();
103                                let existing = decor.prefix().map(|s| s.as_str().unwrap_or("")).unwrap_or("");
104                                let has_existing = existing.trim().starts_with('#');
105
106                                let new_prefix = match self.config.existing_comments {
107                                    ExistingCommentBehavior::Skip if has_existing => None,
108                                    ExistingCommentBehavior::Prepend if has_existing => {
109                                        Some(format!("{}{}", comment, existing))
110                                    }
111                                    ExistingCommentBehavior::Append if has_existing => {
112                                        Some(format!("{}{}", existing, comment))
113                                    }
114                                    _ => Some(comment), // Replace or no existing comment
115                                };
116
117                                if let Some(prefix) = new_prefix {
118                                    decor.set_prefix(prefix);
119                                }
120                            }
121                        }
122                    }
123                }
124            }
125        }
126    }
127}
128
129impl Annotator for TomlAnnotator {
130    fn annotate(
131        &self,
132        content: &str,
133        annotations: &AnnotationMap,
134    ) -> Result<String, AnnotatorError> {
135        let mut doc: DocumentMut = content
136            .parse()
137            .map_err(|e| Error::new(AnnotatorErrorKind::Parse).with_source(e))?;
138
139        self.annotate_table(doc.as_table_mut(), &Vec::new(), annotations);
140
141        Ok(doc.to_string())
142    }
143}
144
145#[cfg(test)]
146mod tests {
147    use super::*;
148    use crate::schema::Annotation;
149    use insta::assert_snapshot;
150
151    fn make_annotations(items: &[(&str, Option<&str>, Option<&str>)]) -> AnnotationMap {
152        let mut map = AnnotationMap::new();
153        for (path, title, desc) in items {
154            let mut ann = Annotation::new(*path);
155            if let Some(t) = title {
156                ann = ann.with_title(*t);
157            }
158            if let Some(d) = desc {
159                ann = ann.with_description(*d);
160            }
161            map.insert(ann);
162        }
163        map
164    }
165
166    #[test]
167    fn test_simple_annotation() {
168        let content = "port = 8080\n";
169        let annotations = make_annotations(&[("port", Some("Port"), Some("Server port number"))]);
170
171        let annotator = TomlAnnotator::new(AnnotatorConfig::default());
172        let result = annotator.annotate(content, &annotations).unwrap();
173
174        assert_snapshot!(result);
175    }
176
177    #[test]
178    fn test_nested_table() {
179        let content = r#"[server]
180port = 8080
181host = "localhost"
182"#;
183        let annotations = make_annotations(&[
184            ("server", Some("Server Config"), None),
185            ("server.port", Some("Port"), Some("The port to listen on")),
186            ("server.host", Some("Host"), None),
187        ]);
188
189        let annotator = TomlAnnotator::new(AnnotatorConfig::default());
190        let result = annotator.annotate(content, &annotations).unwrap();
191
192        assert_snapshot!(result);
193    }
194
195    #[test]
196    fn test_title_only() {
197        let content = "name = \"test\"\n";
198        let annotations = make_annotations(&[("name", Some("Name"), Some("Full description"))]);
199
200        let annotator = TomlAnnotator::new(AnnotatorConfig::titles_only());
201        let result = annotator.annotate(content, &annotations).unwrap();
202
203        assert_snapshot!(result);
204    }
205
206    #[test]
207    fn test_description_only() {
208        let content = "name = \"test\"\n";
209        let annotations = make_annotations(&[("name", Some("Name"), Some("Full description"))]);
210
211        let annotator = TomlAnnotator::new(AnnotatorConfig::descriptions_only());
212        let result = annotator.annotate(content, &annotations).unwrap();
213
214        assert_snapshot!(result);
215    }
216
217    #[test]
218    fn test_preserve_existing_comments() {
219        let content = "# Existing comment\nport = 8080\n";
220        let annotations = make_annotations(&[("port", Some("Port"), None)]);
221
222        let annotator = TomlAnnotator::new(AnnotatorConfig::default());
223        let result = annotator.annotate(content, &annotations).unwrap();
224
225        assert_snapshot!(result);
226    }
227
228    #[test]
229    fn test_no_matching_annotations() {
230        let content = "name = \"test\"\nage = 30\n";
231        let annotations = make_annotations(&[("other", Some("Other"), None)]);
232
233        let annotator = TomlAnnotator::new(AnnotatorConfig::default());
234        let result = annotator.annotate(content, &annotations).unwrap();
235
236        assert_snapshot!(result);
237    }
238
239    #[test]
240    fn test_deeply_nested() {
241        let content = r#"[database]
242[database.connection]
243host = "localhost"
244port = 5432
245"#;
246        let annotations = make_annotations(&[
247            ("database", Some("Database"), None),
248            ("database.connection", Some("Connection Settings"), None),
249            ("database.connection.host", Some("Host"), Some("Database server hostname")),
250            ("database.connection.port", Some("Port"), None),
251        ]);
252
253        let annotator = TomlAnnotator::new(AnnotatorConfig::default());
254        let result = annotator.annotate(content, &annotations).unwrap();
255
256        assert_snapshot!(result);
257    }
258
259    #[test]
260    fn test_long_description_wrapping() {
261        let content = "name = \"test\"\n";
262        let long_desc = "This is a very long description that should be wrapped across multiple lines when the max line width is set to a reasonable value";
263        let annotations = make_annotations(&[("name", None, Some(long_desc))]);
264
265        let config = AnnotatorConfig {
266            max_line_width: Some(40),
267            ..Default::default()
268        };
269        let annotator = TomlAnnotator::new(config);
270        let result = annotator.annotate(content, &annotations).unwrap();
271
272        assert_snapshot!(result);
273    }
274
275    #[test]
276    fn test_skip_existing_comments() {
277        let content = "# Existing comment\nport = 8080\nhost = \"localhost\"\n";
278        let annotations = make_annotations(&[
279            ("port", Some("Port"), None),
280            ("host", Some("Host"), None),
281        ]);
282
283        let config = AnnotatorConfig {
284            existing_comments: ExistingCommentBehavior::Skip,
285            ..Default::default()
286        };
287        let annotator = TomlAnnotator::new(config);
288        let result = annotator.annotate(content, &annotations).unwrap();
289
290        // port should keep its existing comment, host should get the annotation
291        assert_snapshot!(result);
292    }
293
294    #[test]
295    fn test_append_to_existing_comments() {
296        let content = "# Existing comment\nport = 8080\n";
297        let annotations = make_annotations(&[("port", Some("Port"), None)]);
298
299        let config = AnnotatorConfig {
300            existing_comments: ExistingCommentBehavior::Append,
301            ..Default::default()
302        };
303        let annotator = TomlAnnotator::new(config);
304        let result = annotator.annotate(content, &annotations).unwrap();
305
306        assert_snapshot!(result);
307    }
308
309    #[test]
310    fn test_replace_existing_comments() {
311        let content = "# Existing comment\nport = 8080\n";
312        let annotations = make_annotations(&[("port", Some("Port"), None)]);
313
314        let config = AnnotatorConfig {
315            existing_comments: ExistingCommentBehavior::Replace,
316            ..Default::default()
317        };
318        let annotator = TomlAnnotator::new(config);
319        let result = annotator.annotate(content, &annotations).unwrap();
320
321        assert_snapshot!(result);
322    }
323
324    #[test]
325    fn test_include_default_value() {
326        let content = "port = 8080\n";
327
328        let mut map = AnnotationMap::new();
329        map.insert(
330            Annotation::new("port")
331                .with_title("Port")
332                .with_description("The port number")
333                .with_default("3000"),
334        );
335
336        let config = AnnotatorConfig {
337            include_default: true,
338            ..Default::default()
339        };
340        let annotator = TomlAnnotator::new(config);
341        let result = annotator.annotate(content, &map).unwrap();
342
343        assert_snapshot!(result);
344    }
345
346    #[test]
347    fn test_default_value_disabled_by_default() {
348        let content = "port = 8080\n";
349
350        let mut map = AnnotationMap::new();
351        map.insert(
352            Annotation::new("port")
353                .with_title("Port")
354                .with_default("3000"),
355        );
356
357        // Default config has include_default = false
358        let annotator = TomlAnnotator::new(AnnotatorConfig::default());
359        let result = annotator.annotate(content, &map).unwrap();
360
361        assert_snapshot!(result);
362    }
363}
jsonschema_annotator/annotator/toml.rs

jsonschema_annotator/annotator/
toml.rs
