Skip to main content

datui_lib/export/
export_modal.rs

1//! Export modal state and focus management.
2
3use crate::CompressionFormat;
4use crate::widgets::text_input::TextInput;
5use polars::prelude::{LazyFrame, PolarsResult};
6
7#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
8pub enum ExportFormat {
9    #[default]
10    Csv,
11    /// CSV with the tab as its delimiter.
12    Tsv,
13    /// CSV with `|` as its delimiter.
14    Psv,
15    Parquet,
16    Json,
17    Ndjson,
18    /// Arrow IPC / Feather v2
19    Ipc,
20    Avro,
21}
22
23impl ExportFormat {
24    pub const ALL: [Self; 8] = [
25        Self::Csv,
26        Self::Tsv,
27        Self::Psv,
28        Self::Parquet,
29        Self::Json,
30        Self::Ndjson,
31        Self::Ipc,
32        Self::Avro,
33    ];
34
35    pub const fn as_str(self) -> &'static str {
36        match self {
37            Self::Csv => "CSV",
38            Self::Tsv => "TSV",
39            Self::Psv => "PSV",
40            Self::Parquet => "Parquet",
41            Self::Json => "JSON",
42            Self::Ndjson => "NDJSON",
43            Self::Ipc => "Arrow",
44            Self::Avro => "Avro",
45        }
46    }
47
48    pub fn extension(self) -> &'static str {
49        match self {
50            Self::Csv => "csv",
51            Self::Tsv => "tsv",
52            Self::Psv => "psv",
53            Self::Parquet => "parquet",
54            Self::Json => "json",
55            Self::Ndjson => "jsonl",
56            Self::Ipc => "arrow",
57            Self::Avro => "avro",
58        }
59    }
60
61    pub fn from_extension(ext: &str) -> Option<Self> {
62        match ext.to_lowercase().as_str() {
63            "csv" => Some(Self::Csv),
64            "tsv" => Some(Self::Tsv),
65            "psv" => Some(Self::Psv),
66            "parquet" => Some(Self::Parquet),
67            "json" => Some(Self::Json),
68            "ndjson" | "jsonl" => Some(Self::Ndjson),
69            "arrow" | "ipc" | "feather" => Some(Self::Ipc),
70            "avro" => Some(Self::Avro),
71            _ => None,
72        }
73    }
74
75    /// Whether the format stores list, array and struct columns as they are.
76    /// The others get them as JSON text (`nested_json`).
77    pub fn holds_nesting(self) -> bool {
78        !self.is_delimited()
79    }
80
81    /// CSV and its presets, which write delimited text.
82    pub fn is_delimited(self) -> bool {
83        matches!(self, Self::Csv | Self::Tsv | Self::Psv)
84    }
85
86    /// The delimiter a preset sets; `None` for CSV, whose delimiter is the user's.
87    pub fn preset_delimiter(self) -> Option<u8> {
88        match self {
89            Self::Tsv => Some(b'\t'),
90            Self::Psv => Some(b'|'),
91            _ => None,
92        }
93    }
94
95    /// `lf` as this format can write it (planned, not run): binary as base64 and dates
96    /// as text for CSV and JSON, nested columns as JSON and durations as ISO 8601 for
97    /// CSV, and types Avro lacks cast to ones it has.
98    pub fn prepare(self, lf: LazyFrame) -> PolarsResult<LazyFrame> {
99        match self {
100            Self::Csv | Self::Tsv | Self::Psv => crate::export::nested_json::lazy_as_json(lf),
101            Self::Json | Self::Ndjson => crate::export::nested_json::lazy_for_json(lf),
102            Self::Avro => crate::export::avro_types::lazy_for_avro(lf),
103            Self::Parquet | Self::Ipc => Ok(lf),
104        }
105    }
106
107    pub fn supports_compression(self) -> bool {
108        self.is_delimited() || matches!(self, Self::Json | Self::Ndjson)
109    }
110
111    /// The format a path's extension names, looking through a trailing compression
112    /// extension so `out.csv.gz` still names CSV. None for a bare or unknown extension.
113    pub fn from_path(path: &str) -> Option<Self> {
114        let path = std::path::Path::new(path);
115        let ext = path.extension()?.to_str()?;
116        if let Some(format) = Self::from_extension(ext) {
117            return Some(format);
118        }
119        if matches!(ext.to_lowercase().as_str(), "gz" | "zst" | "bz2" | "xz") {
120            let stem = path.file_stem()?.to_str()?;
121            return Self::from_extension(stem.rsplit('.').next()?);
122        }
123        None
124    }
125}
126
127/// The export form's fields.
128#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
129pub enum ExportFocus {
130    #[default]
131    FormatSelector,
132    PathInput,
133    CsvDelimiter,
134    CsvIncludeHeader,
135    /// The compression of the chosen format; offered by the formats that take one.
136    Compression,
137    /// A column naming each row's source file. Offered only when the files disagree,
138    /// where it tells a null from an absent cell downstream.
139    SourceFile,
140}
141
142/// The compressions a delimited or JSON export steps through, in order.
143pub const COMPRESSION_OPTIONS: [Option<CompressionFormat>; 5] = [
144    None,
145    Some(CompressionFormat::Gzip),
146    Some(CompressionFormat::Zstd),
147    Some(CompressionFormat::Bzip2),
148    Some(CompressionFormat::Xz),
149];
150
151pub struct ExportModal {
152    pub focus: ExportFocus,
153    pub selected_format: ExportFormat,
154    pub path_input: TextInput,
155    pub csv_delimiter_input: TextInput,
156    pub csv_include_header: bool,
157    /// Add a column naming the file each row came from. See `ExportFocus::SourceFile`.
158    pub source_file: bool,
159    /// Whether this dataset has files to name. Set when the modal opens.
160    pub offer_source_file: bool,
161    /// Whether the view has list, array or struct columns, which a format
162    /// without nesting writes as JSON; the dialog says so. Set when it opens.
163    pub nested_columns: bool,
164    /// Whether a column or struct field of the view has a name Avro does not
165    /// allow; the dialog says those are renamed. Set when it opens.
166    pub avro_renames: bool,
167    pub csv_compression: Option<CompressionFormat>,
168    pub json_compression: Option<CompressionFormat>,
169    pub ndjson_compression: Option<CompressionFormat>,
170    pub history_limit: usize,
171    /// Why the form cannot export, or why its last write failed, said inline on
172    /// its own status line. Cleared by typing in the path.
173    pub path_error: Option<String>,
174    /// The counts the dialog writes, when it was opened from Value Counts. Kept after
175    /// the form closes while the write runs, so a failed one reopens on them.
176    pub(crate) counts: Option<polars::prelude::DataFrame>,
177}
178
179impl ExportModal {
180    pub fn new() -> Self {
181        Self::default()
182    }
183
184    pub fn open(
185        &mut self,
186        default_format: Option<ExportFormat>,
187        history_limit: usize,
188        theme: &crate::config::Theme,
189        file_delimiter: Option<u8>,
190    ) {
191        self.focus = ExportFocus::PathInput;
192        self.history_limit = history_limit;
193        if let Some(format) = default_format {
194            self.selected_format = format;
195        }
196        // Ctrl+P / Ctrl+N recall the paths exported to before.
197        self.path_input = TextInput::new()
198            .with_history("export_path".to_string())
199            .with_history_limit(history_limit)
200            .with_theme(theme);
201        self.path_input.clear();
202        self.csv_delimiter_input = TextInput::new()
203            .with_history_limit(history_limit)
204            .with_theme(theme);
205        // `--delimiter` if the file was read with one, else a comma.
206        let delimiter_char = file_delimiter.unwrap_or(b',');
207        self.csv_delimiter_input
208            .suggest(format!("{}", delimiter_char as char));
209        self.csv_include_header = true;
210        self.source_file = false;
211        self.offer_source_file = false;
212        self.nested_columns = false;
213        self.avro_renames = false;
214        self.csv_compression = None;
215        self.json_compression = None;
216        self.ndjson_compression = None;
217        self.path_error = None;
218        self.counts = None;
219    }
220
221    pub fn close(&mut self) {
222        self.focus = ExportFocus::FormatSelector;
223        self.path_input.clear();
224        self.path_error = None;
225        self.counts = None;
226    }
227
228    /// The dataset the counts were of is left: a write still running exports them, but
229    /// a failure reopens on the view, not on counts of a dataset gone.
230    pub(crate) fn forget_counts(&mut self) {
231        self.counts = None;
232    }
233
234    /// Follow the typed path's extension with the format choice, so `out.csv` never
235    /// receives Parquet bytes. An unknown extension leaves the format alone; runs only
236    /// when the path changes, so a format picked after typing stands.
237    pub fn sync_format_to_path(&mut self) {
238        let value = self.path_input.value().trim().to_string();
239        if let Some(format) = ExportFormat::from_path(&value) {
240            self.selected_format = format;
241            // A trailing compression extension is part of what the path asks for:
242            // `out.csv.gz` left at Compression: None writes plain bytes to a .gz name.
243            if format.supports_compression()
244                && let Some(comp) = CompressionFormat::from_extension(std::path::Path::new(&value))
245            {
246                self.set_compression_for(format, Some(comp));
247            }
248        }
249    }
250
251    /// Offer `stem` with the format's extension as the path: Enter takes it, typing
252    /// replaces it, and stepping the format carries it along.
253    pub fn suggest_path(&mut self, stem: &str) {
254        let path = format!("{stem}.{}", self.selected_format.extension());
255        self.path_input.suggest(path);
256    }
257
258    /// The reverse of [`Self::sync_format_to_path`]: a format picked after typing
259    /// rewrites the path's format extension. A path with no format extension is left
260    /// alone; a compression suffix stays only if the new format supports one.
261    pub fn sync_path_to_format(&mut self) {
262        let suggested = self.path_input.is_suggested();
263        let value = self.path_input.value().trim().to_string();
264        if value.is_empty() || ExportFormat::from_path(&value).is_none() {
265            return;
266        }
267        let path = std::path::Path::new(&value);
268        let compression = CompressionFormat::from_extension(path)
269            .filter(|_| self.selected_format.supports_compression());
270        // Strip the compression suffix, then the format extension, textually:
271        // Path::set_extension would also eat the `v2` of `data.v2`.
272        let mut base = value.as_str();
273        if CompressionFormat::from_extension(path).is_some()
274            && let Some((rest, _)) = base.rsplit_once('.')
275        {
276            base = rest;
277        }
278        if let Some((rest, ext)) = base.rsplit_once('.')
279            && ExportFormat::from_extension(ext).is_some()
280        {
281            base = rest;
282        }
283        let new_path = match compression {
284            Some(comp) => format!(
285                "{base}.{}.{}",
286                self.selected_format.extension(),
287                comp.extension()
288            ),
289            None => format!("{base}.{}", self.selected_format.extension()),
290        };
291        if suggested {
292            self.path_input.suggest(new_path);
293        } else {
294            self.path_input.set_value(new_path);
295        }
296    }
297
298    /// Set the compression field the given format reads at export time.
299    fn set_compression_for(&mut self, format: ExportFormat, comp: Option<CompressionFormat>) {
300        match format {
301            ExportFormat::Csv | ExportFormat::Tsv | ExportFormat::Psv => {
302                self.csv_compression = comp
303            }
304            ExportFormat::Json => self.json_compression = comp,
305            ExportFormat::Ndjson => self.ndjson_compression = comp,
306            ExportFormat::Parquet | ExportFormat::Ipc | ExportFormat::Avro => {}
307        }
308    }
309
310    /// The compression the chosen format writes with; `None` for one that takes none.
311    pub fn compression(&self) -> Option<CompressionFormat> {
312        match self.selected_format {
313            ExportFormat::Csv | ExportFormat::Tsv | ExportFormat::Psv => self.csv_compression,
314            ExportFormat::Json => self.json_compression,
315            ExportFormat::Ndjson => self.ndjson_compression,
316            ExportFormat::Parquet | ExportFormat::Ipc | ExportFormat::Avro => None,
317        }
318    }
319
320    /// Step the chosen format's compression through [`COMPRESSION_OPTIONS`].
321    pub fn step_compression(&mut self, delta: i8) {
322        let next = crate::app::form::step_value(&COMPRESSION_OPTIONS, self.compression(), delta);
323        self.set_compression_for(self.selected_format, next);
324    }
325
326    /// Step the format, carrying the typed path's extension with it.
327    pub fn step_format(&mut self, delta: i8) {
328        self.selected_format =
329            crate::app::form::step_value(&ExportFormat::ALL, self.selected_format, delta);
330        self.sync_path_to_format();
331        // A format without a delimiter row or compression takes focus off it.
332        crate::app::form::Form::settle_focus(self);
333    }
334
335    /// The fields in Tab order: a list, since options vary by format and dataset.
336    pub fn focus_order(&self) -> Vec<ExportFocus> {
337        crate::app::form::Form::fields(self)
338            .into_iter()
339            .map(|(field, _)| field)
340            .collect()
341    }
342}
343
344impl crate::app::form::Form for ExportModal {
345    type Field = ExportFocus;
346
347    fn fields(&self) -> Vec<(ExportFocus, crate::app::form::FieldKind)> {
348        use crate::app::form::FieldKind::{Checkbox, Choice, Text};
349        let mut fields = vec![
350            (ExportFocus::FormatSelector, Choice),
351            (ExportFocus::PathInput, Text),
352        ];
353        if self.selected_format == ExportFormat::Csv {
354            // The presets say the delimiter.
355            fields.push((ExportFocus::CsvDelimiter, Text));
356        }
357        if self.selected_format.is_delimited() {
358            fields.push((ExportFocus::CsvIncludeHeader, Checkbox));
359        }
360        if self.selected_format.supports_compression() {
361            fields.push((ExportFocus::Compression, Choice));
362        }
363        if self.offer_source_file {
364            fields.push((ExportFocus::SourceFile, Checkbox));
365        }
366        fields
367    }
368
369    fn focused(&self) -> ExportFocus {
370        self.focus
371    }
372
373    fn set_focused(&mut self, field: ExportFocus) {
374        self.focus = field;
375        self.path_input.set_focused(field == ExportFocus::PathInput);
376        self.csv_delimiter_input
377            .set_focused(field == ExportFocus::CsvDelimiter);
378    }
379}
380
381impl Default for ExportModal {
382    fn default() -> Self {
383        Self {
384            focus: ExportFocus::FormatSelector,
385            selected_format: ExportFormat::Csv,
386            path_input: TextInput::new(),
387            csv_delimiter_input: TextInput::new(),
388            csv_include_header: true,
389            source_file: false,
390            offer_source_file: false,
391            nested_columns: false,
392            avro_renames: false,
393            csv_compression: None,
394            json_compression: None,
395            ndjson_compression: None,
396            history_limit: 1000,
397            path_error: None,
398            counts: None,
399        }
400    }
401}
402
403#[cfg(test)]
404mod tests;