use polars::prelude::*;
use serde_json::Value;
use std::sync::Arc;
pub const JSON_INLINE_BYTES: usize = 64 * 1024;
pub const JSON_MAX_BYTES: usize = 4 * 1024 * 1024;
const MEASURED: usize = 1000;
#[derive(Debug, Clone)]
pub enum Node {
Native(Series),
Json { root: Arc<Value>, path: Vec<Step> },
}
#[derive(Debug, Clone)]
pub enum Step {
Key(Arc<str>),
Index(usize),
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Shape {
Struct,
List,
Object,
Array,
Leaf,
}
impl Shape {
pub fn items_title(self) -> &'static str {
match self {
Shape::Struct => "Fields",
Shape::Object => "Keys",
_ => "Items",
}
}
}
impl Node {
pub fn json(&self) -> Option<&Value> {
let Node::Json { root, path } = self else {
return None;
};
let mut at: &Value = root;
for step in path {
at = match (at, step) {
(Value::Object(map), Step::Key(key)) => map.get(&**key)?,
(Value::Array(items), Step::Index(i)) => items.get(*i)?,
_ => return None,
};
}
Some(at)
}
fn is_null(&self) -> bool {
match self {
Node::Native(s) => s.null_count() == s.len(),
Node::Json { .. } => matches!(self.json(), None | Some(Value::Null)),
}
}
pub fn shape(&self) -> Shape {
if self.is_null() {
return Shape::Leaf;
}
match self {
Node::Native(s) => match s.dtype() {
DataType::Struct(_) => Shape::Struct,
DataType::List(_) | DataType::Array(..) => Shape::List,
_ => Shape::Leaf,
},
Node::Json { .. } => match self.json() {
Some(Value::Object(_)) => Shape::Object,
Some(Value::Array(_)) => Shape::Array,
_ => Shape::Leaf,
},
}
}
fn items(&self) -> Option<Series> {
let Node::Native(s) = self else {
return None;
};
match s.get(0).ok()? {
AnyValue::List(inner) | AnyValue::Array(inner, _) => Some(inner),
_ => None,
}
}
fn fields(&self) -> Option<Vec<Series>> {
let Node::Native(s) = self else {
return None;
};
Some(s.struct_().ok()?.fields_as_series())
}
pub fn len(&self) -> usize {
match self.shape() {
Shape::Struct => self.fields().map_or(0, |f| f.len()),
Shape::List => self.items().map_or(0, |s| s.len()),
Shape::Object | Shape::Array => match self.json() {
Some(Value::Object(map)) => map.len(),
Some(Value::Array(items)) => items.len(),
_ => 0,
},
Shape::Leaf => 0,
}
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
fn json_child(&self, step: Step) -> Node {
let Node::Json { root, path } = self else {
unreachable!("only a JSON node has JSON children")
};
let mut path = path.clone();
path.push(step);
Node::Json {
root: Arc::clone(root),
path,
}
}
pub fn children(&self, start: usize, count: usize) -> Vec<(String, Node)> {
match self.shape() {
Shape::Struct => self
.fields()
.unwrap_or_default()
.into_iter()
.skip(start)
.take(count)
.map(|s| (s.name().to_string(), Node::Native(s)))
.collect(),
Shape::List => {
let Some(items) = self.items() else {
return Vec::new();
};
let end = items.len().min(start.saturating_add(count));
(start..end)
.map(|i| (format!("[{i}]"), Node::Native(items.slice(i as i64, 1))))
.collect()
}
Shape::Object => match self.json() {
Some(Value::Object(map)) => page(map.keys(), start, count)
.into_iter()
.map(|key| (key.clone(), self.json_child(Step::Key(key.as_str().into()))))
.collect(),
_ => Vec::new(),
},
Shape::Array => {
let end = self.len().min(start.saturating_add(count));
(start..end)
.map(|i| (format!("[{i}]"), self.json_child(Step::Index(i))))
.collect()
}
Shape::Leaf => Vec::new(),
}
}
pub fn child(&self, i: usize) -> Option<(String, Node)> {
self.children(i, 1).pop()
}
pub fn label_width(&self) -> usize {
match self.shape() {
Shape::List | Shape::Array => format!("[{}]", self.len().saturating_sub(1)).len(),
Shape::Struct => self
.fields()
.unwrap_or_default()
.iter()
.map(|s| crate::glyphs::cell_width(s.name()))
.max()
.unwrap_or(0),
Shape::Object => match self.json() {
Some(Value::Object(map)) => map
.keys()
.take(MEASURED)
.map(|k| crate::glyphs::cell_width(k))
.max()
.unwrap_or(0),
_ => 0,
},
Shape::Leaf => 0,
}
}
pub fn with_text<R>(&self, f: impl FnOnce(&str) -> R) -> Option<R> {
match self {
Node::Native(s) => match s.get(0).ok()? {
AnyValue::String(text) => Some(f(text)),
AnyValue::StringOwned(text) => Some(f(text.as_str())),
_ => None,
},
Node::Json { .. } => match self.json()? {
Value::String(text) => Some(f(text)),
_ => None,
},
}
}
pub fn opens(&self) -> bool {
self.shape() != Shape::Leaf || self.with_text(opens_as_json).unwrap_or(false)
}
pub fn type_label(&self) -> String {
match self {
Node::Native(s) => crate::formats::column_types::dtype_label(s.dtype()),
Node::Json { .. } => json_kind(self.json().unwrap_or(&Value::Null)).to_string(),
}
}
pub fn color_dtype(&self) -> DataType {
match self {
Node::Native(s) => s.dtype().clone(),
Node::Json { .. } => json_dtype(self.json().unwrap_or(&Value::Null)),
}
}
pub fn table_columns(&self) -> Option<Vec<(String, DataType)>> {
match self.shape() {
Shape::List => match self.items()?.dtype() {
DataType::Struct(fields) if !fields.is_empty() => Some(
fields
.iter()
.map(|f| (f.name().to_string(), f.dtype().clone()))
.collect(),
),
_ => None,
},
Shape::Array => match self.json()? {
Value::Array(items) => match items.first()? {
Value::Object(first) if !first.is_empty() => Some(
first
.iter()
.map(|(k, v)| (k.clone(), json_dtype(v)))
.collect(),
),
_ => None,
},
_ => None,
},
_ => None,
}
}
pub fn cell(&self, column: &str) -> Option<Node> {
match self.shape() {
Shape::Struct => match self {
Node::Native(s) => s
.struct_()
.ok()?
.field_by_name(column)
.ok()
.map(Node::Native),
Node::Json { .. } => None,
},
Shape::Object => match self.json()? {
Value::Object(map) if map.contains_key(column) => {
Some(self.json_child(Step::Key(column.into())))
}
_ => None,
},
_ => None,
}
}
}
fn page<T>(
items: impl DoubleEndedIterator<Item = T> + ExactSizeIterator,
start: usize,
count: usize,
) -> Vec<T> {
let len = items.len();
let end = len.min(start.saturating_add(count));
if start >= end {
return Vec::new();
}
if start <= len - end {
return items.skip(start).take(end - start).collect();
}
let mut back: Vec<T> = items.rev().skip(len - end).take(end - start).collect();
back.reverse();
back
}
fn json_dtype(value: &Value) -> DataType {
match value {
Value::String(_) => DataType::String,
Value::Bool(_) => DataType::Boolean,
Value::Number(n) if n.is_f64() => DataType::Float64,
Value::Number(_) => DataType::Int64,
Value::Array(_) => DataType::List(Box::new(DataType::Null)),
Value::Object(_) => DataType::Struct(Vec::new()),
Value::Null => DataType::Null,
}
}
pub fn json_kind(value: &Value) -> &'static str {
match value {
Value::Null => "null",
Value::Bool(_) => "bool",
Value::Number(_) => "number",
Value::String(_) => "str",
Value::Array(_) => "array",
Value::Object(_) => "object",
}
}
pub fn looks_like_json(text: &str) -> bool {
let t = text.trim();
matches!(
(t.as_bytes().first(), t.as_bytes().last()),
(Some(b'{'), Some(b'}')) | (Some(b'['), Some(b']'))
)
}
pub fn opens_as_json(text: &str) -> bool {
text.len() <= JSON_MAX_BYTES && looks_like_json(text)
}
pub fn parse_json(text: &str) -> Result<Value, String> {
if text.len() > JSON_MAX_BYTES {
return Err(format!(
"{} MiB of text is over the {} MiB that opens as JSON",
text.len().div_ceil(1024 * 1024),
JSON_MAX_BYTES / (1024 * 1024)
));
}
serde_json::from_str(text).map_err(|e| format!("not JSON: {e}"))
}
struct Capped {
out: Vec<u8>,
cap: usize,
cut: bool,
}
impl std::io::Write for Capped {
fn write(&mut self, bytes: &[u8]) -> std::io::Result<usize> {
let room = self.cap.saturating_sub(self.out.len());
if bytes.len() > room {
self.out.extend_from_slice(&bytes[..room]);
self.cut = true;
return Err(std::io::Error::other("cap reached"));
}
self.out.extend_from_slice(bytes);
Ok(bytes.len())
}
fn flush(&mut self) -> std::io::Result<()> {
Ok(())
}
}
struct Spaced;
impl serde_json::ser::Formatter for Spaced {
fn begin_array_value<W: ?Sized + std::io::Write>(
&mut self,
writer: &mut W,
first: bool,
) -> std::io::Result<()> {
if first {
Ok(())
} else {
writer.write_all(b", ")
}
}
fn begin_object_key<W: ?Sized + std::io::Write>(
&mut self,
writer: &mut W,
first: bool,
) -> std::io::Result<()> {
if first {
Ok(())
} else {
writer.write_all(b", ")
}
}
fn begin_object_value<W: ?Sized + std::io::Write>(
&mut self,
writer: &mut W,
) -> std::io::Result<()> {
writer.write_all(b": ")
}
}
pub fn json_text(value: &Value, pretty: bool, cap: usize) -> (String, bool) {
use serde::Serialize;
let mut w = Capped {
out: Vec::new(),
cap,
cut: false,
};
let _ = if pretty {
serde_json::to_writer_pretty(&mut w, value)
} else {
value.serialize(&mut serde_json::Serializer::with_formatter(&mut w, Spaced))
};
let text = match String::from_utf8(w.out) {
Ok(text) => text,
Err(e) => {
let valid = e.utf8_error().valid_up_to();
let mut bytes = e.into_bytes();
bytes.truncate(valid);
String::from_utf8(bytes).unwrap_or_default()
}
};
(text, w.cut)
}
pub fn json_copy_text(value: &Value, cap: usize) -> Option<String> {
match value {
Value::String(s) => (s.len() <= cap).then(|| s.clone()),
Value::Null => Some(String::new()),
v => {
let (text, cut) = json_text(v, true, cap);
(!cut).then_some(text)
}
}
}
#[derive(Debug, Clone)]
pub struct Level {
pub label: String,
pub node: Node,
pub selected: usize,
}
impl Level {
pub fn focused(&self) -> Option<(String, Node)> {
self.node.child(self.selected)
}
}
#[derive(Debug, Clone)]
pub struct Drill {
pub frame: u64,
pub row: usize,
pub levels: Vec<Level>,
}
impl Drill {
pub fn level(&self) -> &Level {
self.levels.last().expect("a drill has a level")
}
pub fn level_mut(&mut self) -> &mut Level {
self.levels.last_mut().expect("a drill has a level")
}
}
#[derive(Debug, Clone)]
pub struct JsonWait {
pub token: u64,
pub frame: u64,
pub row: usize,
pub label: String,
pub path: String,
}
pub fn path_key<'a>(steps: impl IntoIterator<Item = &'a str>) -> String {
steps.into_iter().collect::<Vec<_>>().join("\u{1f}")
}
impl Drill {
pub fn item_key(&self, label: &str) -> String {
path_key(
self.levels
.iter()
.map(|l| l.label.as_str())
.chain(std::iter::once(label)),
)
}
}
#[cfg(test)]
mod tests;