use std::{collections::HashSet, fmt::Write, io::IsTerminal};
use inspect_core::{Children, Kind, Sensitivity, ValueRef};
use crate::{color_choice::ColorChoice, role::StyleRole, theme::Theme};
#[derive(Debug, Clone)]
pub struct TreeConfig {
pub show_types: bool,
pub indent: String,
pub use_unicode: bool,
pub redact_secrets: bool,
pub show_truncation: bool,
pub max_items: Option<usize>,
pub max_depth: Option<usize>,
pub color_choice: ColorChoice,
pub theme: Theme,
}
impl Default for TreeConfig {
fn default() -> Self {
Self {
show_types: true,
indent: " ".to_string(),
use_unicode: true,
redact_secrets: true,
show_truncation: true,
max_items: None,
max_depth: None,
color_choice: ColorChoice::Auto,
theme: Theme::default(),
}
}
}
impl TreeConfig {
pub fn compact() -> Self {
Self { show_types: false, show_truncation: false, ..Default::default() }
}
pub fn plain() -> Self {
Self { color_choice: ColorChoice::Never, ..Default::default() }
}
pub fn colored() -> Self {
Self { color_choice: ColorChoice::Always, ..Default::default() }
}
pub fn with_theme(mut self, theme: Theme) -> Self {
self.theme = theme;
self
}
pub fn with_color_choice(mut self, choice: ColorChoice) -> Self {
self.color_choice = choice;
self
}
pub fn with_color(mut self, enabled: bool) -> Self {
self.color_choice = if enabled { ColorChoice::Always } else { ColorChoice::Never };
self
}
pub fn with_max_items(mut self, max: usize) -> Self {
self.max_items = Some(max);
self
}
pub fn with_max_depth(mut self, depth: usize) -> Self {
self.max_depth = Some(depth);
self
}
}
struct StyledWriter<'w, W> {
writer: &'w mut W,
color_enabled: bool,
}
impl<'w, W: Write> StyledWriter<'w, W> {
#[inline(always)]
fn write_plain(&mut self, s: &str) -> std::fmt::Result {
self.writer.write_str(s)
}
#[inline(always)]
fn write_styled(&mut self, text: &str, role: StyleRole, theme: &Theme) -> std::fmt::Result {
if !self.color_enabled {
self.writer.write_str(text)
} else {
let style = theme.style(role);
write!(self.writer, "{}{}{}", style.render(), text, style.render_reset())
}
}
#[inline(always)]
fn write_styled_fmt(
&mut self,
args: std::fmt::Arguments<'_>,
role: StyleRole,
theme: &Theme,
) -> std::fmt::Result {
if !self.color_enabled {
self.writer.write_fmt(args)
} else {
let style = theme.style(role);
write!(self.writer, "{}", style.render())?;
self.writer.write_fmt(args)?;
write!(self.writer, "{}", style.render_reset())
}
}
}
pub struct TreeFormatter {
config: TreeConfig,
visited: HashSet<usize>,
}
impl TreeFormatter {
pub fn new() -> Self {
Self::with_config(TreeConfig::default())
}
pub fn with_config(config: TreeConfig) -> Self {
Self { config, visited: HashSet::new() }
}
pub fn format(&mut self, value: &ValueRef<'_>) -> String {
let is_term = std::io::stdout().is_terminal();
let color_enabled = self.config.color_choice.should_color(is_term);
let mut output = String::new();
self.format_to_styled(value, &mut output, color_enabled).ok();
output
}
pub fn format_colored(&mut self, value: &ValueRef<'_>) -> String {
let mut output = String::new();
self.format_to_styled(value, &mut output, true).ok();
output
}
pub fn format_plain(&mut self, value: &ValueRef<'_>) -> String {
let mut output = String::new();
self.format_to_styled(value, &mut output, false).ok();
output
}
pub fn format_to<W: Write>(
&mut self,
value: &ValueRef<'_>,
writer: &mut W,
) -> std::fmt::Result {
let is_term = std::io::stdout().is_terminal();
let color_enabled = self.config.color_choice.should_color(is_term);
self.format_to_styled(value, writer, color_enabled)
}
pub fn format_to_styled<W: Write>(
&mut self,
value: &ValueRef<'_>,
writer: &mut W,
color_enabled: bool,
) -> std::fmt::Result {
self.visited.clear();
let mut styled_writer = StyledWriter { writer, color_enabled };
self.format_value(value, &mut styled_writer, "", true, 0)
}
pub fn format_io<W: std::io::Write>(
&mut self,
value: &ValueRef<'_>,
mut writer: W,
) -> std::io::Result<()> {
let mut buf = String::new();
self.format_to(value, &mut buf).map_err(std::io::Error::other)?;
writer.write_all(buf.as_bytes())
}
fn format_value<W: Write>(
&mut self,
value: &ValueRef<'_>,
out: &mut StyledWriter<'_, W>,
prefix: &str,
is_last: bool,
current_depth: usize,
) -> std::fmt::Result {
if value.sensitivity() == Sensitivity::Hidden {
return Ok(());
}
let (branch, continuation) = if prefix.is_empty() {
("", "")
} else if self.config.use_unicode {
if is_last { ("└── ", " ") } else { ("├── ", "│ ") }
} else {
if is_last { ("+-- ", " ") } else { ("|-- ", "| ") }
};
if !prefix.is_empty() {
out.write_styled(prefix, StyleRole::Punctuation, &self.config.theme)?;
out.write_styled(branch, StyleRole::Punctuation, &self.config.theme)?;
}
if value.sensitivity() == Sensitivity::Secret && self.config.redact_secrets {
out.write_styled("[REDACTED]", StyleRole::Sensitive, &self.config.theme)?;
out.write_plain("\n")?;
return Ok(());
}
if let Some(max_depth) = self.config.max_depth {
if current_depth >= max_depth {
out.write_styled(
"… <max depth reached>",
StyleRole::Truncated,
&self.config.theme,
)?;
out.write_plain("\n")?;
return Ok(());
}
}
match value.kind() {
Kind::Unit => {
out.write_styled("()", StyleRole::Null, &self.config.theme)?;
}
Kind::Bool(b) => {
out.write_styled_fmt(format_args!("{b}"), StyleRole::Boolean, &self.config.theme)?;
}
Kind::Char(c) => {
out.write_styled("'", StyleRole::Punctuation, &self.config.theme)?;
let mut esc_buf = String::new();
for esc in c.escape_debug() {
esc_buf.push(esc);
}
out.write_styled(&esc_buf, StyleRole::String, &self.config.theme)?;
out.write_styled("'", StyleRole::Punctuation, &self.config.theme)?;
}
Kind::I8(n) => {
out.write_styled_fmt(format_args!("{n}"), StyleRole::Number, &self.config.theme)?
}
Kind::I16(n) => {
out.write_styled_fmt(format_args!("{n}"), StyleRole::Number, &self.config.theme)?
}
Kind::I32(n) => {
out.write_styled_fmt(format_args!("{n}"), StyleRole::Number, &self.config.theme)?
}
Kind::I64(n) => {
out.write_styled_fmt(format_args!("{n}"), StyleRole::Number, &self.config.theme)?
}
Kind::I128(n) => {
out.write_styled_fmt(format_args!("{n}"), StyleRole::Number, &self.config.theme)?
}
Kind::Isize(n) => {
out.write_styled_fmt(format_args!("{n}"), StyleRole::Number, &self.config.theme)?
}
Kind::U8(n) => {
out.write_styled_fmt(format_args!("{n}"), StyleRole::Number, &self.config.theme)?
}
Kind::U16(n) => {
out.write_styled_fmt(format_args!("{n}"), StyleRole::Number, &self.config.theme)?
}
Kind::U32(n) => {
out.write_styled_fmt(format_args!("{n}"), StyleRole::Number, &self.config.theme)?
}
Kind::U64(n) => {
out.write_styled_fmt(format_args!("{n}"), StyleRole::Number, &self.config.theme)?
}
Kind::U128(n) => {
out.write_styled_fmt(format_args!("{n}"), StyleRole::Number, &self.config.theme)?
}
Kind::Usize(n) => {
out.write_styled_fmt(format_args!("{n}"), StyleRole::Number, &self.config.theme)?
}
Kind::F32(f) => {
out.write_styled_fmt(format_args!("{f}"), StyleRole::Number, &self.config.theme)?
}
Kind::F64(f) => {
out.write_styled_fmt(format_args!("{f}"), StyleRole::Number, &self.config.theme)?
}
Kind::Str(s) => {
out.write_styled("\"", StyleRole::Punctuation, &self.config.theme)?;
let mut esc_buf = String::new();
for c in s.chars() {
for esc in c.escape_debug() {
esc_buf.push(esc);
}
}
out.write_styled(&esc_buf, StyleRole::String, &self.config.theme)?;
out.write_styled("\"", StyleRole::Punctuation, &self.config.theme)?;
}
Kind::Bytes(b) => {
out.write_styled("[", StyleRole::Punctuation, &self.config.theme)?;
out.write_styled_fmt(
format_args!("{} bytes", b.len()),
StyleRole::Metadata,
&self.config.theme,
)?;
out.write_styled("]", StyleRole::Punctuation, &self.config.theme)?;
}
_ => {
let type_name = value.type_info().name();
if self.config.show_types {
out.write_styled(type_name, StyleRole::Type, &self.config.theme)?;
}
if let Some(variant) = value.variant() {
if self.config.show_types {
out.write_styled("::", StyleRole::Punctuation, &self.config.theme)?;
}
out.write_styled(variant.name(), StyleRole::Variant, &self.config.theme)?;
}
}
}
out.write_plain("\n")?;
if let Some(children) = value.children() {
let next_prefix = if prefix.is_empty() {
"".to_string()
} else {
format!("{}{}", prefix, continuation)
};
self.format_children(children, out, &next_prefix, value.kind(), current_depth + 1)?;
}
Ok(())
}
fn format_children<W: Write>(
&mut self,
children: &Children<'_>,
out: &mut StyledWriter<'_, W>,
prefix: &str,
parent_kind: &Kind<'_>,
current_depth: usize,
) -> std::fmt::Result {
match children {
Children::Direct(fields) => {
let total_count = fields.len();
let limit = self.config.max_items.unwrap_or(total_count).min(total_count);
let is_truncated = limit < total_count;
for (idx, (field, child_value)) in fields.iter().take(limit).enumerate() {
let is_last = (idx == limit - 1) && !is_truncated;
let (branch, continuation) = if self.config.use_unicode {
if is_last { ("└── ", " ") } else { ("├── ", "│ ") }
} else {
if is_last { ("+-- ", " ") } else { ("|-- ", "| ") }
};
if field.sensitivity() == Sensitivity::Hidden
|| child_value.sensitivity() == Sensitivity::Hidden
{
continue;
}
out.write_styled(prefix, StyleRole::Punctuation, &self.config.theme)?;
out.write_styled(branch, StyleRole::Punctuation, &self.config.theme)?;
if let Some(name) = field.name() {
let role = if matches!(parent_kind, Kind::Map) {
StyleRole::Key
} else {
StyleRole::Field
};
out.write_styled(name, role, &self.config.theme)?;
out.write_styled(": ", StyleRole::Punctuation, &self.config.theme)?;
} else {
out.write_styled("[", StyleRole::Punctuation, &self.config.theme)?;
out.write_styled_fmt(
format_args!("{}", field.index()),
StyleRole::Index,
&self.config.theme,
)?;
out.write_styled("]: ", StyleRole::Punctuation, &self.config.theme)?;
}
if (field.sensitivity() == Sensitivity::Secret
|| child_value.sensitivity() == Sensitivity::Secret)
&& self.config.redact_secrets
{
out.write_styled("[REDACTED]", StyleRole::Sensitive, &self.config.theme)?;
out.write_plain("\n")?;
continue;
}
if child_value.kind().is_scalar() {
self.format_scalar_inline(child_value, out)?;
out.write_plain("\n")?;
} else {
if let Some(max_depth) = self.config.max_depth {
if current_depth >= max_depth {
out.write_styled(
"… <max depth reached>",
StyleRole::Truncated,
&self.config.theme,
)?;
out.write_plain("\n")?;
continue;
}
}
let type_name = child_value.type_info().name();
if self.config.show_types {
out.write_styled(type_name, StyleRole::Type, &self.config.theme)?;
}
if let Some(variant) = child_value.variant() {
if self.config.show_types {
out.write_styled("::", StyleRole::Punctuation, &self.config.theme)?;
}
out.write_styled(
variant.name(),
StyleRole::Variant,
&self.config.theme,
)?;
}
out.write_plain("\n")?;
if let Some(sub_children) = child_value.children() {
let next_prefix = format!("{}{}", prefix, continuation);
self.format_children(
sub_children,
out,
&next_prefix,
child_value.kind(),
current_depth + 1,
)?;
}
}
}
if is_truncated && self.config.show_truncation {
let remaining = total_count - limit;
let branch = if self.config.use_unicode { "└── " } else { "\\-- " };
out.write_styled(prefix, StyleRole::Punctuation, &self.config.theme)?;
out.write_styled(branch, StyleRole::Punctuation, &self.config.theme)?;
let notice = format!("… {remaining} more");
out.write_styled(¬ice, StyleRole::Truncated, &self.config.theme)?;
out.write_plain("\n")?;
}
}
}
Ok(())
}
fn format_scalar_inline<W: Write>(
&self,
value: &ValueRef<'_>,
out: &mut StyledWriter<'_, W>,
) -> std::fmt::Result {
match value.kind() {
Kind::Unit => out.write_styled("()", StyleRole::Null, &self.config.theme),
Kind::Bool(b) => {
out.write_styled_fmt(format_args!("{b}"), StyleRole::Boolean, &self.config.theme)
}
Kind::Char(c) => {
out.write_styled("'", StyleRole::Punctuation, &self.config.theme)?;
let mut esc_buf = String::new();
for esc in c.escape_debug() {
esc_buf.push(esc);
}
out.write_styled(&esc_buf, StyleRole::String, &self.config.theme)?;
out.write_styled("'", StyleRole::Punctuation, &self.config.theme)
}
Kind::I8(n) => {
out.write_styled_fmt(format_args!("{n}"), StyleRole::Number, &self.config.theme)
}
Kind::I16(n) => {
out.write_styled_fmt(format_args!("{n}"), StyleRole::Number, &self.config.theme)
}
Kind::I32(n) => {
out.write_styled_fmt(format_args!("{n}"), StyleRole::Number, &self.config.theme)
}
Kind::I64(n) => {
out.write_styled_fmt(format_args!("{n}"), StyleRole::Number, &self.config.theme)
}
Kind::I128(n) => {
out.write_styled_fmt(format_args!("{n}"), StyleRole::Number, &self.config.theme)
}
Kind::Isize(n) => {
out.write_styled_fmt(format_args!("{n}"), StyleRole::Number, &self.config.theme)
}
Kind::U8(n) => {
out.write_styled_fmt(format_args!("{n}"), StyleRole::Number, &self.config.theme)
}
Kind::U16(n) => {
out.write_styled_fmt(format_args!("{n}"), StyleRole::Number, &self.config.theme)
}
Kind::U32(n) => {
out.write_styled_fmt(format_args!("{n}"), StyleRole::Number, &self.config.theme)
}
Kind::U64(n) => {
out.write_styled_fmt(format_args!("{n}"), StyleRole::Number, &self.config.theme)
}
Kind::U128(n) => {
out.write_styled_fmt(format_args!("{n}"), StyleRole::Number, &self.config.theme)
}
Kind::Usize(n) => {
out.write_styled_fmt(format_args!("{n}"), StyleRole::Number, &self.config.theme)
}
Kind::F32(f) => {
out.write_styled_fmt(format_args!("{f}"), StyleRole::Number, &self.config.theme)
}
Kind::F64(f) => {
out.write_styled_fmt(format_args!("{f}"), StyleRole::Number, &self.config.theme)
}
Kind::Str(s) => {
out.write_styled("\"", StyleRole::Punctuation, &self.config.theme)?;
let mut esc_buf = String::new();
for c in s.chars() {
for esc in c.escape_debug() {
esc_buf.push(esc);
}
}
out.write_styled(&esc_buf, StyleRole::String, &self.config.theme)?;
out.write_styled("\"", StyleRole::Punctuation, &self.config.theme)
}
Kind::Bytes(b) => {
out.write_styled("[", StyleRole::Punctuation, &self.config.theme)?;
out.write_styled_fmt(
format_args!("{} bytes", b.len()),
StyleRole::Metadata,
&self.config.theme,
)?;
out.write_styled("]", StyleRole::Punctuation, &self.config.theme)
}
_ => out.write_styled(value.type_info().name(), StyleRole::Type, &self.config.theme),
}
}
}
impl Default for TreeFormatter {
fn default() -> Self {
Self::new()
}
}