jsonschema_annotator/annotator/
toml.rs

1use toml_edit::{DocumentMut, Item, Table};
2
3use super::{Annotator, AnnotatorConfig};
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 lines.is_empty() {
36            None
37        } else {
38            // Add newline after comments so it appears before the key
39            Some(lines.join("\n") + "\n")
40        }
41    }
42
43    fn annotate_table(
44        &self,
45        table: &mut Table,
46        path: &[String],
47        annotations: &AnnotationMap,
48    ) {
49        // Collect keys first to avoid borrow issues
50        // Use deref to str to get the key string (Key derefs to str)
51        let keys: Vec<String> = table.iter().map(|(k, _)| (*k).to_string()).collect();
52
53        for key_string in keys {
54            let mut current_path = path.to_vec();
55            current_path.push(key_string.clone());
56            let path_string = current_path.join(".");
57
58            // Get mutable access to the key-value pair
59            if let Some((mut key, item)) = table.get_key_value_mut(&key_string) {
60                // Handle tables vs regular values differently
61                match item {
62                    Item::Table(nested) => {
63                        // For tables, use the table's own decor (appears before the [header])
64                        if let Some(ann) = annotations.get(&path_string) {
65                            if let Some(comment) = self.format_comment(ann) {
66                                let decor = nested.decor_mut();
67                                if self.config.preserve_existing {
68                                    let existing = decor.prefix().map(|s| s.as_str().unwrap_or("")).unwrap_or("");
69                                    if !existing.is_empty() {
70                                        decor.set_prefix(format!("{}{}", comment, existing));
71                                    } else {
72                                        decor.set_prefix(comment);
73                                    }
74                                } else {
75                                    decor.set_prefix(comment);
76                                }
77                            }
78                        }
79                        // Recurse into nested tables
80                        self.annotate_table(nested, &current_path, annotations);
81                    }
82                    Item::Value(toml_edit::Value::InlineTable(_)) => {
83                        // Can't easily modify inline tables, skip for now
84                    }
85                    _ => {
86                        // For regular values, use the key's decor
87                        if let Some(ann) = annotations.get(&path_string) {
88                            if let Some(comment) = self.format_comment(ann) {
89                                let decor = key.leaf_decor_mut();
90                                if self.config.preserve_existing {
91                                    let existing = decor.prefix().map(|s| s.as_str().unwrap_or("")).unwrap_or("");
92                                    if !existing.is_empty() {
93                                        decor.set_prefix(format!("{}{}", comment, existing));
94                                    } else {
95                                        decor.set_prefix(comment);
96                                    }
97                                } else {
98                                    decor.set_prefix(comment);
99                                }
100                            }
101                        }
102                    }
103                }
104            }
105        }
106    }
107}
108
109impl Annotator for TomlAnnotator {
110    fn annotate(
111        &self,
112        content: &str,
113        annotations: &AnnotationMap,
114    ) -> Result<String, AnnotatorError> {
115        let mut doc: DocumentMut = content
116            .parse()
117            .map_err(|e| Error::new(AnnotatorErrorKind::Parse).with_source(e))?;
118
119        self.annotate_table(doc.as_table_mut(), &Vec::new(), annotations);
120
121        Ok(doc.to_string())
122    }
123}
124
125#[cfg(test)]
126mod tests {
127    use super::*;
128    use crate::schema::Annotation;
129    use insta::assert_snapshot;
130
131    fn make_annotations(items: &[(&str, Option<&str>, Option<&str>)]) -> AnnotationMap {
132        let mut map = AnnotationMap::new();
133        for (path, title, desc) in items {
134            let mut ann = Annotation::new(*path);
135            if let Some(t) = title {
136                ann = ann.with_title(*t);
137            }
138            if let Some(d) = desc {
139                ann = ann.with_description(*d);
140            }
141            map.insert(ann);
142        }
143        map
144    }
145
146    #[test]
147    fn test_simple_annotation() {
148        let content = "port = 8080\n";
149        let annotations = make_annotations(&[("port", Some("Port"), Some("Server port number"))]);
150
151        let annotator = TomlAnnotator::new(AnnotatorConfig::default());
152        let result = annotator.annotate(content, &annotations).unwrap();
153
154        assert_snapshot!(result);
155    }
156
157    #[test]
158    fn test_nested_table() {
159        let content = r#"[server]
160port = 8080
161host = "localhost"
162"#;
163        let annotations = make_annotations(&[
164            ("server", Some("Server Config"), None),
165            ("server.port", Some("Port"), Some("The port to listen on")),
166            ("server.host", Some("Host"), None),
167        ]);
168
169        let annotator = TomlAnnotator::new(AnnotatorConfig::default());
170        let result = annotator.annotate(content, &annotations).unwrap();
171
172        assert_snapshot!(result);
173    }
174
175    #[test]
176    fn test_title_only() {
177        let content = "name = \"test\"\n";
178        let annotations = make_annotations(&[("name", Some("Name"), Some("Full description"))]);
179
180        let annotator = TomlAnnotator::new(AnnotatorConfig::titles_only());
181        let result = annotator.annotate(content, &annotations).unwrap();
182
183        assert_snapshot!(result);
184    }
185
186    #[test]
187    fn test_description_only() {
188        let content = "name = \"test\"\n";
189        let annotations = make_annotations(&[("name", Some("Name"), Some("Full description"))]);
190
191        let annotator = TomlAnnotator::new(AnnotatorConfig::descriptions_only());
192        let result = annotator.annotate(content, &annotations).unwrap();
193
194        assert_snapshot!(result);
195    }
196
197    #[test]
198    fn test_preserve_existing_comments() {
199        let content = "# Existing comment\nport = 8080\n";
200        let annotations = make_annotations(&[("port", Some("Port"), None)]);
201
202        let annotator = TomlAnnotator::new(AnnotatorConfig::default());
203        let result = annotator.annotate(content, &annotations).unwrap();
204
205        assert_snapshot!(result);
206    }
207
208    #[test]
209    fn test_no_matching_annotations() {
210        let content = "name = \"test\"\nage = 30\n";
211        let annotations = make_annotations(&[("other", Some("Other"), None)]);
212
213        let annotator = TomlAnnotator::new(AnnotatorConfig::default());
214        let result = annotator.annotate(content, &annotations).unwrap();
215
216        assert_snapshot!(result);
217    }
218
219    #[test]
220    fn test_deeply_nested() {
221        let content = r#"[database]
222[database.connection]
223host = "localhost"
224port = 5432
225"#;
226        let annotations = make_annotations(&[
227            ("database", Some("Database"), None),
228            ("database.connection", Some("Connection Settings"), None),
229            ("database.connection.host", Some("Host"), Some("Database server hostname")),
230            ("database.connection.port", Some("Port"), None),
231        ]);
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_long_description_wrapping() {
241        let content = "name = \"test\"\n";
242        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";
243        let annotations = make_annotations(&[("name", None, Some(long_desc))]);
244
245        let mut config = AnnotatorConfig::default();
246        config.max_line_width = Some(40);
247        let annotator = TomlAnnotator::new(config);
248        let result = annotator.annotate(content, &annotations).unwrap();
249
250        assert_snapshot!(result);
251    }
252}
jsonschema_annotator/annotator/toml.rs

jsonschema_annotator/annotator/
toml.rs
