1use polars::prelude::*;
10use serde_json::Value;
11use std::sync::Arc;
12
13pub const JSON_INLINE_BYTES: usize = 64 * 1024;
16pub const JSON_MAX_BYTES: usize = 4 * 1024 * 1024;
18const MEASURED: usize = 1000;
21
22#[derive(Debug, Clone)]
24pub enum Node {
25 Native(Series),
27 Json { root: Arc<Value>, path: Vec<Step> },
29}
30
31#[derive(Debug, Clone)]
35pub enum Step {
36 Key(Arc<str>),
37 Index(usize),
38}
39
40#[derive(Debug, Clone, Copy, PartialEq, Eq)]
42pub enum Shape {
43 Struct,
44 List,
45 Object,
46 Array,
47 Leaf,
49}
50
51impl Shape {
52 pub fn items_title(self) -> &'static str {
54 match self {
55 Shape::Struct => "Fields",
56 Shape::Object => "Keys",
57 _ => "Items",
58 }
59 }
60}
61
62impl Node {
63 pub fn json(&self) -> Option<&Value> {
65 let Node::Json { root, path } = self else {
66 return None;
67 };
68 let mut at: &Value = root;
69 for step in path {
70 at = match (at, step) {
71 (Value::Object(map), Step::Key(key)) => map.get(&**key)?,
72 (Value::Array(items), Step::Index(i)) => items.get(*i)?,
73 _ => return None,
74 };
75 }
76 Some(at)
77 }
78
79 fn is_null(&self) -> bool {
80 match self {
81 Node::Native(s) => s.null_count() == s.len(),
82 Node::Json { .. } => matches!(self.json(), None | Some(Value::Null)),
83 }
84 }
85
86 pub fn shape(&self) -> Shape {
87 if self.is_null() {
88 return Shape::Leaf;
89 }
90 match self {
91 Node::Native(s) => match s.dtype() {
92 DataType::Struct(_) => Shape::Struct,
93 DataType::List(_) | DataType::Array(..) => Shape::List,
94 _ => Shape::Leaf,
95 },
96 Node::Json { .. } => match self.json() {
97 Some(Value::Object(_)) => Shape::Object,
98 Some(Value::Array(_)) => Shape::Array,
99 _ => Shape::Leaf,
100 },
101 }
102 }
103
104 fn items(&self) -> Option<Series> {
106 let Node::Native(s) = self else {
107 return None;
108 };
109 match s.get(0).ok()? {
110 AnyValue::List(inner) | AnyValue::Array(inner, _) => Some(inner),
111 _ => None,
112 }
113 }
114
115 fn fields(&self) -> Option<Vec<Series>> {
117 let Node::Native(s) = self else {
118 return None;
119 };
120 Some(s.struct_().ok()?.fields_as_series())
121 }
122
123 pub fn len(&self) -> usize {
125 match self.shape() {
126 Shape::Struct => self.fields().map_or(0, |f| f.len()),
127 Shape::List => self.items().map_or(0, |s| s.len()),
128 Shape::Object | Shape::Array => match self.json() {
129 Some(Value::Object(map)) => map.len(),
130 Some(Value::Array(items)) => items.len(),
131 _ => 0,
132 },
133 Shape::Leaf => 0,
134 }
135 }
136
137 pub fn is_empty(&self) -> bool {
138 self.len() == 0
139 }
140
141 fn json_child(&self, step: Step) -> Node {
142 let Node::Json { root, path } = self else {
143 unreachable!("only a JSON node has JSON children")
144 };
145 let mut path = path.clone();
146 path.push(step);
147 Node::Json {
148 root: Arc::clone(root),
149 path,
150 }
151 }
152
153 pub fn children(&self, start: usize, count: usize) -> Vec<(String, Node)> {
156 match self.shape() {
157 Shape::Struct => self
158 .fields()
159 .unwrap_or_default()
160 .into_iter()
161 .skip(start)
162 .take(count)
163 .map(|s| (s.name().to_string(), Node::Native(s)))
164 .collect(),
165 Shape::List => {
166 let Some(items) = self.items() else {
167 return Vec::new();
168 };
169 let end = items.len().min(start.saturating_add(count));
170 (start..end)
171 .map(|i| (format!("[{i}]"), Node::Native(items.slice(i as i64, 1))))
172 .collect()
173 }
174 Shape::Object => match self.json() {
175 Some(Value::Object(map)) => page(map.keys(), start, count)
176 .into_iter()
177 .map(|key| (key.clone(), self.json_child(Step::Key(key.as_str().into()))))
178 .collect(),
179 _ => Vec::new(),
180 },
181 Shape::Array => {
182 let end = self.len().min(start.saturating_add(count));
183 (start..end)
184 .map(|i| (format!("[{i}]"), self.json_child(Step::Index(i))))
185 .collect()
186 }
187 Shape::Leaf => Vec::new(),
188 }
189 }
190
191 pub fn child(&self, i: usize) -> Option<(String, Node)> {
192 self.children(i, 1).pop()
193 }
194
195 pub fn label_width(&self) -> usize {
198 match self.shape() {
199 Shape::List | Shape::Array => format!("[{}]", self.len().saturating_sub(1)).len(),
200 Shape::Struct => self
201 .fields()
202 .unwrap_or_default()
203 .iter()
204 .map(|s| crate::glyphs::cell_width(s.name()))
205 .max()
206 .unwrap_or(0),
207 Shape::Object => match self.json() {
208 Some(Value::Object(map)) => map
209 .keys()
210 .take(MEASURED)
211 .map(|k| crate::glyphs::cell_width(k))
212 .max()
213 .unwrap_or(0),
214 _ => 0,
215 },
216 Shape::Leaf => 0,
217 }
218 }
219
220 pub fn with_text<R>(&self, f: impl FnOnce(&str) -> R) -> Option<R> {
222 match self {
223 Node::Native(s) => match s.get(0).ok()? {
224 AnyValue::String(text) => Some(f(text)),
225 AnyValue::StringOwned(text) => Some(f(text.as_str())),
226 _ => None,
227 },
228 Node::Json { .. } => match self.json()? {
229 Value::String(text) => Some(f(text)),
230 _ => None,
231 },
232 }
233 }
234
235 pub fn opens(&self) -> bool {
238 self.shape() != Shape::Leaf || self.with_text(opens_as_json).unwrap_or(false)
239 }
240
241 pub fn type_label(&self) -> String {
243 match self {
244 Node::Native(s) => crate::widgets::datatable::dtype_label(s.dtype()),
245 Node::Json { .. } => json_kind(self.json().unwrap_or(&Value::Null)).to_string(),
246 }
247 }
248
249 pub fn color_dtype(&self) -> DataType {
251 match self {
252 Node::Native(s) => s.dtype().clone(),
253 Node::Json { .. } => json_dtype(self.json().unwrap_or(&Value::Null)),
254 }
255 }
256
257 pub fn table_columns(&self) -> Option<Vec<(String, DataType)>> {
261 match self.shape() {
262 Shape::List => match self.items()?.dtype() {
263 DataType::Struct(fields) if !fields.is_empty() => Some(
264 fields
265 .iter()
266 .map(|f| (f.name().to_string(), f.dtype().clone()))
267 .collect(),
268 ),
269 _ => None,
270 },
271 Shape::Array => match self.json()? {
272 Value::Array(items) => match items.first()? {
273 Value::Object(first) if !first.is_empty() => Some(
274 first
275 .iter()
276 .map(|(k, v)| (k.clone(), json_dtype(v)))
277 .collect(),
278 ),
279 _ => None,
280 },
281 _ => None,
282 },
283 _ => None,
284 }
285 }
286
287 pub fn cell(&self, column: &str) -> Option<Node> {
290 match self.shape() {
291 Shape::Struct => match self {
293 Node::Native(s) => s
294 .struct_()
295 .ok()?
296 .field_by_name(column)
297 .ok()
298 .map(Node::Native),
299 Node::Json { .. } => None,
300 },
301 Shape::Object => match self.json()? {
302 Value::Object(map) if map.contains_key(column) => {
303 Some(self.json_child(Step::Key(column.into())))
304 }
305 _ => None,
306 },
307 _ => None,
308 }
309 }
310}
311
312fn page<T>(
316 items: impl DoubleEndedIterator<Item = T> + ExactSizeIterator,
317 start: usize,
318 count: usize,
319) -> Vec<T> {
320 let len = items.len();
321 let end = len.min(start.saturating_add(count));
322 if start >= end {
323 return Vec::new();
324 }
325 if start <= len - end {
326 return items.skip(start).take(end - start).collect();
327 }
328 let mut back: Vec<T> = items.rev().skip(len - end).take(end - start).collect();
329 back.reverse();
330 back
331}
332
333fn json_dtype(value: &Value) -> DataType {
335 match value {
336 Value::String(_) => DataType::String,
337 Value::Bool(_) => DataType::Boolean,
338 Value::Number(n) if n.is_f64() => DataType::Float64,
339 Value::Number(_) => DataType::Int64,
340 Value::Array(_) => DataType::List(Box::new(DataType::Null)),
341 Value::Object(_) => DataType::Struct(Vec::new()),
342 Value::Null => DataType::Null,
343 }
344}
345
346pub fn json_kind(value: &Value) -> &'static str {
348 match value {
349 Value::Null => "null",
350 Value::Bool(_) => "bool",
351 Value::Number(_) => "number",
352 Value::String(_) => "str",
353 Value::Array(_) => "array",
354 Value::Object(_) => "object",
355 }
356}
357
358pub fn looks_like_json(text: &str) -> bool {
361 let t = text.trim();
362 matches!(
363 (t.as_bytes().first(), t.as_bytes().last()),
364 (Some(b'{'), Some(b'}')) | (Some(b'['), Some(b']'))
365 )
366}
367
368pub fn opens_as_json(text: &str) -> bool {
371 text.len() <= JSON_MAX_BYTES && looks_like_json(text)
372}
373
374pub fn parse_json(text: &str) -> Result<Value, String> {
377 if text.len() > JSON_MAX_BYTES {
378 return Err(format!(
379 "{} MiB of text is over the {} MiB that opens as JSON",
380 text.len().div_ceil(1024 * 1024),
381 JSON_MAX_BYTES / (1024 * 1024)
382 ));
383 }
384 serde_json::from_str(text).map_err(|e| format!("not JSON: {e}"))
385}
386
387struct Capped {
389 out: Vec<u8>,
390 cap: usize,
391 cut: bool,
392}
393
394impl std::io::Write for Capped {
395 fn write(&mut self, bytes: &[u8]) -> std::io::Result<usize> {
396 let room = self.cap.saturating_sub(self.out.len());
397 if bytes.len() > room {
398 self.out.extend_from_slice(&bytes[..room]);
399 self.cut = true;
400 return Err(std::io::Error::other("cap reached"));
401 }
402 self.out.extend_from_slice(bytes);
403 Ok(bytes.len())
404 }
405
406 fn flush(&mut self) -> std::io::Result<()> {
407 Ok(())
408 }
409}
410
411struct Spaced;
414
415impl serde_json::ser::Formatter for Spaced {
416 fn begin_array_value<W: ?Sized + std::io::Write>(
417 &mut self,
418 writer: &mut W,
419 first: bool,
420 ) -> std::io::Result<()> {
421 if first {
422 Ok(())
423 } else {
424 writer.write_all(b", ")
425 }
426 }
427
428 fn begin_object_key<W: ?Sized + std::io::Write>(
429 &mut self,
430 writer: &mut W,
431 first: bool,
432 ) -> std::io::Result<()> {
433 if first {
434 Ok(())
435 } else {
436 writer.write_all(b", ")
437 }
438 }
439
440 fn begin_object_value<W: ?Sized + std::io::Write>(
441 &mut self,
442 writer: &mut W,
443 ) -> std::io::Result<()> {
444 writer.write_all(b": ")
445 }
446}
447
448pub fn json_text(value: &Value, pretty: bool, cap: usize) -> (String, bool) {
451 use serde::Serialize;
452 let mut w = Capped {
453 out: Vec::new(),
454 cap,
455 cut: false,
456 };
457 let _ = if pretty {
459 serde_json::to_writer_pretty(&mut w, value)
460 } else {
461 value.serialize(&mut serde_json::Serializer::with_formatter(&mut w, Spaced))
462 };
463 let text = match String::from_utf8(w.out) {
464 Ok(text) => text,
465 Err(e) => {
467 let valid = e.utf8_error().valid_up_to();
468 let mut bytes = e.into_bytes();
469 bytes.truncate(valid);
470 String::from_utf8(bytes).unwrap_or_default()
471 }
472 };
473 (text, w.cut)
474}
475
476pub fn json_copy_text(value: &Value, cap: usize) -> Option<String> {
479 match value {
480 Value::String(s) => (s.len() <= cap).then(|| s.clone()),
481 Value::Null => Some(String::new()),
482 v => {
483 let (text, cut) = json_text(v, true, cap);
484 (!cut).then_some(text)
485 }
486 }
487}
488
489#[derive(Debug, Clone)]
491pub struct Level {
492 pub label: String,
494 pub node: Node,
495 pub selected: usize,
496}
497
498impl Level {
499 pub fn focused(&self) -> Option<(String, Node)> {
501 self.node.child(self.selected)
502 }
503}
504
505#[derive(Debug, Clone)]
507pub struct Drill {
508 pub frame: u64,
509 pub row: usize,
510 pub levels: Vec<Level>,
511}
512
513impl Drill {
514 pub fn level(&self) -> &Level {
515 self.levels.last().expect("a drill has a level")
516 }
517
518 pub fn level_mut(&mut self) -> &mut Level {
519 self.levels.last_mut().expect("a drill has a level")
520 }
521}
522
523#[derive(Debug, Clone)]
525pub struct JsonWait {
526 pub token: u64,
527 pub frame: u64,
528 pub row: usize,
529 pub label: String,
530 pub path: String,
532}
533
534pub fn path_key<'a>(steps: impl IntoIterator<Item = &'a str>) -> String {
538 steps.into_iter().collect::<Vec<_>>().join("\u{1f}")
539}
540
541impl Drill {
542 pub fn item_key(&self, label: &str) -> String {
544 path_key(
545 self.levels
546 .iter()
547 .map(|l| l.label.as_str())
548 .chain(std::iter::once(label)),
549 )
550 }
551}
552
553#[cfg(test)]
554mod tests {
555 use super::*;
556
557 fn struct_row() -> Series {
558 let df = df!(
559 "id" => [7i64],
560 "name" => ["ann"],
561 )
562 .unwrap();
563 df.into_struct("customer".into()).into_series()
564 }
565
566 #[test]
567 fn a_struct_lists_its_fields_as_one_row_slices() {
568 let node = Node::Native(struct_row());
569 assert_eq!(node.shape(), Shape::Struct);
570 assert_eq!(node.len(), 2);
571 let kids = node.children(0, 10);
572 let labels: Vec<_> = kids.iter().map(|(l, _)| l.as_str()).collect();
573 assert_eq!(labels, ["id", "name"]);
574 let (_, name) = &kids[1];
575 assert_eq!(name.shape(), Shape::Leaf);
576 assert_eq!(name.with_text(str::to_string).as_deref(), Some("ann"));
577 assert_eq!(node.cell("id").unwrap().type_label(), "i64");
578 assert!(node.cell("nope").is_none());
579 }
580
581 #[test]
583 fn a_long_list_makes_only_the_page_asked_for() {
584 let inner = Series::new("".into(), (0..1_000_000i64).collect::<Vec<_>>());
585 let list = Series::new("l".into(), [inner]);
586 let node = Node::Native(list);
587 assert_eq!(node.shape(), Shape::List);
588 assert_eq!(node.len(), 1_000_000);
589 let page = node.children(999_998, 20);
590 assert_eq!(page.len(), 2);
591 assert_eq!(page[0].0, "[999998]");
592 assert_eq!(page[1].1.type_label(), "i64");
593 assert_eq!(node.label_width(), "[999999]".len());
594 }
595
596 #[test]
597 fn a_null_struct_or_list_is_a_leaf() {
598 let list = Series::new_null("l".into(), 1).cast(&DataType::List(Box::new(DataType::Int64)));
599 let node = Node::Native(list.unwrap());
600 assert_eq!(node.shape(), Shape::Leaf);
601 assert!(!node.opens());
602 }
603
604 #[test]
605 fn json_objects_keep_their_order_and_arrays_their_items() {
606 let root = Arc::new(parse_json(r#"{"z": 1, "a": [true, null, {"k": "v"}]}"#).unwrap());
607 let node = Node::Json {
608 root,
609 path: Vec::new(),
610 };
611 assert_eq!(node.shape(), Shape::Object);
612 let kids = node.children(0, 9);
613 assert_eq!(kids[0].0, "z", "the document's order, not sorted");
614 assert_eq!(kids[0].1.type_label(), "number");
615 let (_, a) = &kids[1];
616 assert_eq!(a.shape(), Shape::Array);
617 assert_eq!(a.len(), 3);
618 let (label, obj) = a.child(2).unwrap();
619 assert_eq!(label, "[2]");
620 assert_eq!(obj.shape(), Shape::Object);
621 assert_eq!(obj.cell("k").unwrap().json(), Some(&Value::from("v")));
622 assert_eq!(a.child(1).unwrap().1.shape(), Shape::Leaf);
623 }
624
625 #[test]
628 fn a_large_objects_last_keys_resolve_by_key() {
629 let n = 200_000;
630 let text = format!(
631 "{{{}}}",
632 (0..n)
633 .map(|i| format!("\"k{i}\": {i}"))
634 .collect::<Vec<_>>()
635 .join(", ")
636 );
637 let node = Node::Json {
638 root: Arc::new(parse_json(&text).unwrap()),
639 path: Vec::new(),
640 };
641 let tail = node.children(n - 3, 10);
642 let labels: Vec<_> = tail.iter().map(|(l, _)| l.as_str()).collect();
643 assert_eq!(labels, ["k199997", "k199998", "k199999"]);
644 assert_eq!(tail[2].1.json(), Some(&Value::from(n - 1)));
645 let Node::Json { path, .. } = &tail[2].1 else {
646 panic!("a JSON item");
647 };
648 assert!(matches!(&path[..], [Step::Key(k)] if &**k == "k199999"));
649 let head = node.children(1, 2);
650 assert_eq!(head[1].0, "k2");
651 assert_eq!(page(0..10, 7, 5), [7, 8, 9]);
652 assert_eq!(page(0..10, 2, 3), [2, 3, 4]);
653 assert!(page(0..10, 10, 3).is_empty());
654 }
655
656 #[test]
657 fn text_opens_only_when_it_reads_as_an_object_or_array() {
658 assert!(looks_like_json(" {\"a\": 1}\n"));
659 assert!(looks_like_json("[1, 2]"));
660 assert!(!looks_like_json("{not closed"));
661 assert!(!looks_like_json("\"text\""));
662 assert!(!looks_like_json("42"));
663 let s = Series::new("s".into(), ["[1, 2]"]);
664 assert!(Node::Native(s).opens());
665 let s = Series::new("s".into(), ["plain"]);
666 assert!(!Node::Native(s).opens());
667 }
668
669 #[test]
672 fn parsing_is_bounded() {
673 let deep = format!("{}{}", "[".repeat(10_000), "]".repeat(10_000));
674 assert!(parse_json(&deep).is_err());
675 let huge = format!("[{}0]", "0,".repeat(JSON_MAX_BYTES / 2));
676 let err = parse_json(&huge).unwrap_err();
677 assert!(err.contains("MiB"), "{err}");
678 assert!(parse_json("{oops}").unwrap_err().starts_with("not JSON"));
679 }
680
681 #[test]
682 fn json_text_stops_at_its_cap() {
683 let value: Value = serde_json::from_str(&format!("[{}1]", "1,".repeat(10_000))).unwrap();
684 let (text, cut) = json_text(&value, true, 100);
685 assert!(cut);
686 assert_eq!(text.len(), 100);
687 let (whole, cut) = json_text(&Value::from("é"), false, 100);
688 assert!(!cut);
689 assert_eq!(whole, "\"é\"");
690 let (text, cut) = json_text(&Value::from("ééé"), false, 4);
692 assert!(cut);
693 assert_eq!(text, "\"é");
694 assert_eq!(json_copy_text(&value, 100), None);
695 assert_eq!(json_copy_text(&Value::from("raw"), 100).unwrap(), "raw");
696 }
697
698 #[test]
699 fn a_list_of_structs_and_an_array_of_objects_are_tables() {
700 let items = df!("sku" => ["A1", "B7"], "qty" => [2i64, 1])
701 .unwrap()
702 .into_struct("".into())
703 .into_series();
704 let list = Series::new("items".into(), [items]);
705 let node = Node::Native(list);
706 assert_eq!(
707 node.table_columns(),
708 Some(vec![
709 ("sku".to_string(), DataType::String),
710 ("qty".to_string(), DataType::Int64)
711 ])
712 );
713 let (_, first) = node.child(0).unwrap();
714 assert_eq!(first.shape(), Shape::Struct);
715 assert_eq!(
716 first
717 .cell("sku")
718 .unwrap()
719 .with_text(str::to_string)
720 .as_deref(),
721 Some("A1")
722 );
723
724 let root = Arc::new(parse_json(r#"[{"a": 1, "b": 2}, {"b": 3}]"#).unwrap());
725 let arr = Node::Json {
726 root,
727 path: Vec::new(),
728 };
729 assert_eq!(
730 arr.table_columns(),
731 Some(vec![
732 ("a".to_string(), DataType::Int64),
733 ("b".to_string(), DataType::Int64)
734 ])
735 );
736 assert!(arr.child(1).unwrap().1.cell("a").is_none());
737 }
738}