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