Skip to main content

vortex_array/display/
extractor.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::fmt;
5
6use crate::ArrayRef;
7
8/// Context threaded through tree traversal for percentage calculations etc.
9pub struct TreeContext {
10    /// Stack of ancestor nbytes values. `None` entries reset the percentage root
11    /// (e.g. for chunked arrays where each chunk is its own root).
12    pub(crate) ancestor_sizes: Vec<Option<u64>>,
13}
14
15impl TreeContext {
16    pub(crate) fn new() -> Self {
17        Self {
18            ancestor_sizes: Vec::new(),
19        }
20    }
21
22    /// The total size used as the denominator for percentage calculations.
23    /// Returns `None` if there is no ancestor (i.e., this node is the root or
24    /// a chunk boundary reset the percentage root).
25    pub fn parent_total_size(&self) -> Option<u64> {
26        self.ancestor_sizes.last().cloned().flatten()
27    }
28
29    pub(crate) fn push(&mut self, size: Option<u64>) {
30        self.ancestor_sizes.push(size);
31    }
32
33    pub(crate) fn pop(&mut self) {
34        self.ancestor_sizes.pop();
35    }
36}
37
38/// Wrapper providing access to a [`fmt::Formatter`] and the current indentation string.
39pub struct IndentedFormatter<'a, 'b> {
40    inner: &'a mut fmt::Formatter<'b>,
41    indent: &'a str,
42}
43
44impl<'a, 'b> IndentedFormatter<'a, 'b> {
45    pub(crate) fn new(f: &'a mut fmt::Formatter<'b>, indent: &'a str) -> Self {
46        Self { inner: f, indent }
47    }
48
49    /// Access the indent string and underlying [`fmt::Formatter`] together.
50    pub fn parts(&mut self) -> (&str, &mut fmt::Formatter<'b>) {
51        (self.indent, self.inner)
52    }
53
54    /// The current indentation string.
55    pub fn indent(&self) -> &str {
56        self.indent
57    }
58
59    /// Access the underlying [`fmt::Formatter`].
60    pub fn formatter(&mut self) -> &mut fmt::Formatter<'b> {
61        self.inner
62    }
63}
64
65/// Trait for contributing display information to tree nodes.
66///
67/// Each extractor represents one "dimension" of display (e.g., nbytes, stats, metadata, buffers).
68/// Extractors are composable: you can combine any number of them via [`TreeDisplay::with`].
69///
70/// [`TreeDisplay::with`]: super::TreeDisplay::with
71pub trait TreeExtractor: Send + Sync {
72    /// Write header annotations (space-prefixed) to the formatter.
73    fn write_header(
74        &self,
75        array: &ArrayRef,
76        ctx: &TreeContext,
77        f: &mut fmt::Formatter<'_>,
78    ) -> fmt::Result {
79        let _ = (array, ctx, f);
80        Ok(())
81    }
82
83    /// Write detail lines below the header.
84    ///
85    /// Content written through `f` is automatically indented. Use
86    /// [`f.formatter()`](IndentedFormatter::formatter) to access the underlying
87    /// [`fmt::Formatter`] for formatting flags.
88    fn write_details(
89        &self,
90        array: &ArrayRef,
91        ctx: &TreeContext,
92        f: &mut IndentedFormatter<'_, '_>,
93    ) -> fmt::Result {
94        let _ = (array, ctx, f);
95        Ok(())
96    }
97}