Skip to main content

rich_ext/data/
transform.rs

1//! Transforms over a data [`Document`]: redact, select, filter and highlight
2//! by JSONPath (the `jsonpath` feature). See [`crate::transform`] for the
3//! pipeline.
4//!
5//! ```
6//! use rich::{Console, Style};
7//! use rich_ext::data::transform::{Document, Filter, Select};
8//! use rich_ext::data::{parse, Format};
9//! use rich_ext::transform::Pipeline;
10//!
11//! let node = parse(Format::Json, r#"{"a": {"b": 1, "c": 2}, "d": [1, 2]}"#).unwrap();
12//! let pipeline = Pipeline::new().then("filter", Filter::new("$.a.c").unwrap());
13//! let document = pipeline.apply(Document::new(node).label("doc")).unwrap();
14//! let out = Console::builder().width(40).build().render_export(&document.explorer());
15//! assert_eq!(out, "doc\n└── a\n    └── c: 2\n");
16//! ```
17
18use std::collections::HashSet;
19
20use rich::Style;
21
22use super::{Explorer, Node, Path, Redactor, SelectError, Selectors, Value};
23use crate::transform::{Transform, TransformError};
24
25/// A data tree on its way to being rendered, with what transforms decided
26/// about how to show it.
27#[derive(Clone, Debug, PartialEq)]
28pub struct Document {
29    pub node: Node,
30    /// The root's label: the file name, or the path [`Select`] narrowed to.
31    pub label: Option<String>,
32    /// Styles for the lines of these paths, from [`Highlight`].
33    pub highlights: Vec<(Path, Style)>,
34}
35
36impl Document {
37    pub fn new(node: Node) -> Self {
38        Document {
39            node,
40            label: None,
41            highlights: Vec::new(),
42        }
43    }
44
45    #[must_use]
46    pub fn label(mut self, label: impl Into<String>) -> Self {
47        self.label = Some(label.into());
48        self
49    }
50
51    /// An [`Explorer`] over the document, with its label and highlights.
52    pub fn explorer(&self) -> Explorer<'_> {
53        let mut explorer = Explorer::new(&self.node);
54        if let Some(label) = &self.label {
55            explorer = explorer.root_label(label.clone());
56        }
57        for (path, style) in &self.highlights {
58            explorer = explorer.highlight(path.clone(), style.clone());
59        }
60        explorer
61    }
62}
63
64/// Masks values with a [`Redactor`], such as
65/// [`Redaction::secrets`](super::Redaction::secrets).
66pub struct Redact<R>(pub R);
67
68impl<R: Redactor + Send + Sync> Transform<Document> for Redact<R> {
69    fn apply(&self, mut document: Document) -> Result<Document, TransformError> {
70        document.node = document.node.redacted(&self.0);
71        Ok(document)
72    }
73}
74
75/// A compiled JSONPath expression, kept as text so the transform is `Sync`.
76#[derive(Clone, Debug)]
77struct Expression(String);
78
79impl Expression {
80    fn new(expression: &str) -> Result<Self, SelectError> {
81        Selectors::default().compile("jsonpath", expression)?;
82        Ok(Expression(expression.to_string()))
83    }
84
85    fn hits<'a>(&self, node: &'a Node) -> Result<Vec<(Path, &'a Node)>, TransformError> {
86        let to_error = |e: SelectError| TransformError::new(e.to_string());
87        Selectors::default()
88            .compile("jsonpath", &self.0)
89            .map_err(to_error)?
90            .select(node)
91            .map_err(to_error)
92    }
93}
94
95/// Narrows a document to what a JSONPath selects. One hit becomes the root,
96/// labelled with its path; several (or none) become a map keyed by path.
97#[derive(Clone, Debug)]
98pub struct Select(Expression);
99
100impl Select {
101    pub fn new(expression: &str) -> Result<Self, SelectError> {
102        Expression::new(expression).map(Select)
103    }
104}
105
106impl Transform<Document> for Select {
107    fn apply(&self, mut document: Document) -> Result<Document, TransformError> {
108        let hits: Vec<_> = self
109            .0
110            .hits(&document.node)?
111            .into_iter()
112            .map(|(path, hit)| (path.to_string(), hit.clone()))
113            .collect();
114        document.node = match <[_; 1]>::try_from(hits) {
115            Ok([(path, hit)]) => {
116                document.label = Some(if path.is_empty() { "$".into() } else { path });
117                hit
118            }
119            Err(hits) => Node::new(Value::Map(hits)),
120        };
121        // Paths no longer name the same nodes.
122        document.highlights.clear();
123        Ok(document)
124    }
125}
126
127/// Keeps what a JSONPath selects and the containers above it, and drops the
128/// rest, so the document keeps its shape. Sequences are renumbered.
129#[derive(Clone, Debug)]
130pub struct Filter(Expression);
131
132impl Filter {
133    pub fn new(expression: &str) -> Result<Self, SelectError> {
134        Expression::new(expression).map(Filter)
135    }
136}
137
138impl Transform<Document> for Filter {
139    fn apply(&self, mut document: Document) -> Result<Document, TransformError> {
140        let keep: HashSet<Path> = self
141            .0
142            .hits(&document.node)?
143            .into_iter()
144            .map(|(path, _)| path)
145            .collect();
146        let root = &document.node;
147        document.node = prune(root, &Path::root(), &keep).unwrap_or_else(|| {
148            let empty = match root.value {
149                Value::Seq(_) => Value::Seq(Vec::new()),
150                _ => Value::Map(Vec::new()),
151            };
152            Node {
153                value: empty,
154                meta: root.meta.clone(),
155            }
156        });
157        document.highlights.clear();
158        Ok(document)
159    }
160}
161
162fn prune(node: &Node, path: &Path, keep: &HashSet<Path>) -> Option<Node> {
163    if keep.contains(path) {
164        return Some(node.clone());
165    }
166    let value = match &node.value {
167        Value::Seq(items) => {
168            let kept: Vec<Node> = items
169                .iter()
170                .enumerate()
171                .filter_map(|(i, item)| prune(item, &path.child_index(i), keep))
172                .collect();
173            (!kept.is_empty()).then_some(Value::Seq(kept))?
174        }
175        Value::Map(entries) => {
176            let kept: Vec<(String, Node)> = entries
177                .iter()
178                .filter_map(|(key, item)| {
179                    prune(item, &path.child_key(key), keep).map(|item| (key.clone(), item))
180                })
181                .collect();
182            (!kept.is_empty()).then_some(Value::Map(kept))?
183        }
184        _ => return None,
185    };
186    Some(Node {
187        value,
188        meta: node.meta.clone(),
189    })
190}
191
192/// Styles the tree lines of what a JSONPath selects.
193#[derive(Clone, Debug)]
194pub struct Highlight {
195    expression: Expression,
196    style: Style,
197}
198
199impl Highlight {
200    pub fn new(expression: &str, style: Style) -> Result<Self, SelectError> {
201        Ok(Highlight {
202            expression: Expression::new(expression)?,
203            style,
204        })
205    }
206}
207
208impl Transform<Document> for Highlight {
209    fn apply(&self, mut document: Document) -> Result<Document, TransformError> {
210        let paths: Vec<Path> = self
211            .expression
212            .hits(&document.node)?
213            .into_iter()
214            .map(|(path, _)| path)
215            .collect();
216        document
217            .highlights
218            .extend(paths.into_iter().map(|path| (path, self.style.clone())));
219        Ok(document)
220    }
221}
222
223#[cfg(test)]
224mod tests {
225    use super::*;
226    use crate::data::{Format, Redaction, parse};
227    use crate::transform::Pipeline;
228
229    fn doc(json: &str) -> Document {
230        Document::new(parse(Format::Json, json).unwrap())
231    }
232
233    #[test]
234    fn select_narrows_and_labels() {
235        let one = Select::new("$.a.b")
236            .unwrap()
237            .apply(doc(r#"{"a": {"b": [1]}}"#).label("file"))
238            .unwrap();
239        assert_eq!(one.label.as_deref(), Some("a.b"));
240        assert_eq!(one.node.to_json(), serde_json::json!([1]));
241
242        let root = Select::new("$").unwrap().apply(doc("[1]")).unwrap();
243        assert_eq!(root.label.as_deref(), Some("$"));
244
245        let many = Select::new("$..x")
246            .unwrap()
247            .apply(doc(r#"{"x": 1, "y": {"x": 2}}"#).label("file"))
248            .unwrap();
249        assert_eq!(many.label.as_deref(), Some("file"));
250        assert_eq!(many.node.to_json(), serde_json::json!({"x": 1, "y.x": 2}));
251        assert!(Select::new("$[").is_err());
252    }
253
254    #[test]
255    fn filter_keeps_shape() {
256        let filtered = Filter::new("$..name")
257            .unwrap()
258            .apply(doc(
259                r#"{"a": [{"name": "x", "age": 1}, {"age": 2}], "b": 3}"#,
260            ))
261            .unwrap();
262        assert_eq!(
263            filtered.node.to_json(),
264            serde_json::json!({"a": [{"name": "x"}]})
265        );
266
267        let none = Filter::new("$.missing")
268            .unwrap()
269            .apply(doc("[1, 2]"))
270            .unwrap();
271        assert_eq!(none.node.to_json(), serde_json::json!([]));
272    }
273
274    #[test]
275    fn highlight_styles_selected_lines() {
276        let pipeline = Pipeline::new()
277            .then("redact", Redact(Redaction::secrets()))
278            .then(
279                "highlight",
280                Highlight::new("$.b", Style::parse("reverse").unwrap()).unwrap(),
281            );
282        let document = pipeline
283            .apply(doc(r#"{"a": 1, "b": 2, "password": "x"}"#).label("doc"))
284            .unwrap();
285        assert_eq!(document.highlights.len(), 1);
286        let console = rich::Console::builder()
287            .width(40)
288            .force_terminal(true)
289            .color_system(Some(rich::ColorSystem::Standard))
290            .build();
291        let out = console.render_to_string(&document.explorer());
292        let line = out.lines().find(|l| l.contains('b')).unwrap();
293        assert!(line.contains("\x1b[7m"), "{line:?}");
294        assert!(
295            !out.lines()
296                .find(|l| l.contains("a"))
297                .unwrap()
298                .contains("\x1b[7m")
299        );
300        assert!(out.contains("********"));
301    }
302}