Skip to main content

datui_lib/
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: binary as base64 and dates as text for
96    /// CSV and JSON, nested columns as JSON and durations as ISO 8601 for CSV, and the types
97    /// Avro lacks cast to ones it has. Planned, not run.
98    pub fn prepare(self, lf: LazyFrame) -> PolarsResult<LazyFrame> {
99        match self {
100            Self::Csv | Self::Tsv | Self::Psv => crate::nested_json::lazy_as_json(lf),
101            Self::Json | Self::Ndjson => crate::nested_json::lazy_for_json(lf),
102            Self::Avro => crate::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    /// Add a column naming the file each row came from. Only offered for a dataset
138    /// whose files disagree, since that is where a null and an absent cell differ and
139    /// the source file is what tells them apart downstream.
140    SourceFile,
141}
142
143/// The compressions a delimited or JSON export steps through, in order.
144pub const COMPRESSION_OPTIONS: [Option<CompressionFormat>; 5] = [
145    None,
146    Some(CompressionFormat::Gzip),
147    Some(CompressionFormat::Zstd),
148    Some(CompressionFormat::Bzip2),
149    Some(CompressionFormat::Xz),
150];
151
152pub struct ExportModal {
153    pub active: bool,
154    pub focus: ExportFocus,
155    pub selected_format: ExportFormat,
156    pub path_input: TextInput,
157    // CSV options
158    pub csv_delimiter_input: TextInput,
159    pub csv_include_header: bool,
160    /// Add a column naming the file each row came from. See `ExportFocus::SourceFile`.
161    pub source_file: bool,
162    /// Whether this dataset has files to name. Set when the modal opens.
163    pub offer_source_file: bool,
164    /// Whether the view has list, array or struct columns, which a format
165    /// without nesting writes as JSON; the dialog says so. Set when it opens.
166    pub nested_columns: bool,
167    /// Whether a column or struct field of the view has a name Avro does not
168    /// allow; the dialog says those are renamed. Set when it opens.
169    pub avro_renames: bool,
170    pub csv_compression: Option<CompressionFormat>,
171    // JSON options
172    pub json_compression: Option<CompressionFormat>,
173    // NDJSON options
174    pub ndjson_compression: Option<CompressionFormat>,
175    pub history_limit: usize,
176    /// Why the form cannot export, or why its last write failed, said inline on
177    /// its own status line. Cleared by typing in the path.
178    pub path_error: Option<String>,
179}
180
181impl ExportModal {
182    pub fn new() -> Self {
183        Self::default()
184    }
185
186    pub fn open(
187        &mut self,
188        default_format: Option<ExportFormat>,
189        history_limit: usize,
190        theme: &crate::config::Theme,
191        file_delimiter: Option<u8>,
192    ) {
193        self.active = true;
194        self.focus = ExportFocus::PathInput;
195        self.history_limit = history_limit;
196        if let Some(format) = default_format {
197            self.selected_format = format;
198        }
199        // Ctrl+P / Ctrl+N recall the paths exported to before.
200        self.path_input = TextInput::new()
201            .with_history("export_path".to_string())
202            .with_history_limit(history_limit)
203            .with_theme(theme);
204        self.path_input.clear();
205        self.csv_delimiter_input = TextInput::new()
206            .with_history_limit(history_limit)
207            .with_theme(theme);
208        // `--delimiter` if the file was read with one, else a comma.
209        let delimiter_char = file_delimiter.unwrap_or(b',');
210        self.csv_delimiter_input
211            .suggest(format!("{}", delimiter_char as char));
212        self.csv_include_header = true;
213        self.source_file = false;
214        self.offer_source_file = false;
215        self.nested_columns = false;
216        self.avro_renames = false;
217        self.csv_compression = None;
218        self.json_compression = None;
219        self.ndjson_compression = None;
220        self.path_error = None;
221    }
222
223    pub fn close(&mut self) {
224        self.active = false;
225        self.focus = ExportFocus::FormatSelector;
226        self.path_input.clear();
227        self.path_error = None;
228    }
229
230    /// Hide behind a child confirmation without discarding the form; `resume`
231    /// brings it back exactly as typed. `close` is the discard.
232    pub fn suspend(&mut self) {
233        self.active = false;
234    }
235
236    pub fn resume(&mut self) {
237        self.active = true;
238    }
239
240    /// Follow the typed path's extension with the format picker, so `out.csv` never
241    /// silently receives Parquet bytes. An extension that names no format leaves the
242    /// picker alone, and an explicit format picked after typing stands, because this
243    /// runs only when the path itself changes.
244    pub fn sync_format_to_path(&mut self) {
245        let value = self.path_input.value().trim().to_string();
246        if let Some(format) = ExportFormat::from_path(&value) {
247            self.selected_format = format;
248            // A trailing compression extension is part of what the path asks for:
249            // `out.csv.gz` left at Compression: None writes plain bytes to a .gz name.
250            if format.supports_compression()
251                && let Some(comp) = CompressionFormat::from_extension(std::path::Path::new(&value))
252            {
253                self.set_compression_for(format, Some(comp));
254            }
255        }
256    }
257
258    /// The other direction of [`Self::sync_format_to_path`]: a format picked after
259    /// typing rewrites the path's format extension, so `out.csv` never silently
260    /// receives Parquet bytes from the picker side either. A path whose extension
261    /// names no format is left alone. A compression suffix survives when the new
262    /// format supports one and is dropped when it cannot.
263    /// Offer `stem` with the format's extension as the path: Enter takes it as it
264    /// stands, typing replaces it, and stepping the format carries it along.
265    pub fn suggest_path(&mut self, stem: &str) {
266        let path = format!("{stem}.{}", self.selected_format.extension());
267        self.path_input.suggest(path);
268    }
269
270    pub fn sync_path_to_format(&mut self) {
271        let suggested = self.path_input.is_suggested();
272        let value = self.path_input.value().trim().to_string();
273        if value.is_empty() || ExportFormat::from_path(&value).is_none() {
274            return;
275        }
276        let path = std::path::Path::new(&value);
277        let compression = CompressionFormat::from_extension(path)
278            .filter(|_| self.selected_format.supports_compression());
279        // Strip the compression suffix, then the format extension, textually:
280        // Path::set_extension would also eat the `v2` of `data.v2`.
281        let mut base = value.as_str();
282        if CompressionFormat::from_extension(path).is_some()
283            && let Some((rest, _)) = base.rsplit_once('.')
284        {
285            base = rest;
286        }
287        if let Some((rest, ext)) = base.rsplit_once('.')
288            && ExportFormat::from_extension(ext).is_some()
289        {
290            base = rest;
291        }
292        let new_path = match compression {
293            Some(comp) => format!(
294                "{base}.{}.{}",
295                self.selected_format.extension(),
296                comp.extension()
297            ),
298            None => format!("{base}.{}", self.selected_format.extension()),
299        };
300        if suggested {
301            self.path_input.suggest(new_path);
302        } else {
303            self.path_input.set_value(new_path);
304        }
305    }
306
307    /// Set the compression field the given format reads at export time.
308    fn set_compression_for(&mut self, format: ExportFormat, comp: Option<CompressionFormat>) {
309        match format {
310            ExportFormat::Csv | ExportFormat::Tsv | ExportFormat::Psv => {
311                self.csv_compression = comp
312            }
313            ExportFormat::Json => self.json_compression = comp,
314            ExportFormat::Ndjson => self.ndjson_compression = comp,
315            ExportFormat::Parquet | ExportFormat::Ipc | ExportFormat::Avro => {}
316        }
317    }
318
319    /// The compression the chosen format writes with; `None` for one that takes none.
320    pub fn compression(&self) -> Option<CompressionFormat> {
321        match self.selected_format {
322            ExportFormat::Csv | ExportFormat::Tsv | ExportFormat::Psv => self.csv_compression,
323            ExportFormat::Json => self.json_compression,
324            ExportFormat::Ndjson => self.ndjson_compression,
325            ExportFormat::Parquet | ExportFormat::Ipc | ExportFormat::Avro => None,
326        }
327    }
328
329    /// Step the chosen format's compression through [`COMPRESSION_OPTIONS`].
330    pub fn step_compression(&mut self, delta: i8) {
331        let next = crate::form::step_value(&COMPRESSION_OPTIONS, self.compression(), delta);
332        self.set_compression_for(self.selected_format, next);
333    }
334
335    /// Step the format, carrying the typed path's extension with it.
336    pub fn step_format(&mut self, delta: i8) {
337        self.selected_format =
338            crate::form::step_value(&ExportFormat::ALL, self.selected_format, delta);
339        self.sync_path_to_format();
340        // A format without a delimiter row or compression takes focus off it.
341        crate::form::Form::settle_focus(self);
342    }
343
344    /// The fields this modal offers, in the order Tab walks them.
345    ///
346    /// Built as a list rather than a match per field: the options differ by format and
347    /// one of them depends on the dataset, and a hand-written state machine over both
348    /// has an arm for every pair.
349    pub fn focus_order(&self) -> Vec<ExportFocus> {
350        crate::form::Form::fields(self)
351            .into_iter()
352            .map(|(field, _)| field)
353            .collect()
354    }
355}
356
357impl crate::form::Form for ExportModal {
358    type Field = ExportFocus;
359
360    fn fields(&self) -> Vec<(ExportFocus, crate::form::FieldKind)> {
361        use crate::form::FieldKind::{Checkbox, Choice, Text};
362        let mut fields = vec![
363            (ExportFocus::FormatSelector, Choice),
364            (ExportFocus::PathInput, Text),
365        ];
366        if self.selected_format == ExportFormat::Csv {
367            // The presets say the delimiter.
368            fields.push((ExportFocus::CsvDelimiter, Text));
369        }
370        if self.selected_format.is_delimited() {
371            fields.push((ExportFocus::CsvIncludeHeader, Checkbox));
372        }
373        if self.selected_format.supports_compression() {
374            fields.push((ExportFocus::Compression, Choice));
375        }
376        if self.offer_source_file {
377            fields.push((ExportFocus::SourceFile, Checkbox));
378        }
379        fields
380    }
381
382    fn focused(&self) -> ExportFocus {
383        self.focus
384    }
385
386    fn set_focused(&mut self, field: ExportFocus) {
387        self.focus = field;
388        self.path_input.set_focused(field == ExportFocus::PathInput);
389        self.csv_delimiter_input
390            .set_focused(field == ExportFocus::CsvDelimiter);
391    }
392}
393
394impl Default for ExportModal {
395    fn default() -> Self {
396        Self {
397            active: false,
398            focus: ExportFocus::FormatSelector,
399            selected_format: ExportFormat::Csv,
400            path_input: TextInput::new(),
401            csv_delimiter_input: TextInput::new(),
402            csv_include_header: true,
403            source_file: false,
404            offer_source_file: false,
405            nested_columns: false,
406            avro_renames: false,
407            csv_compression: None,
408            json_compression: None,
409            ndjson_compression: None,
410            history_limit: 1000,
411            path_error: None,
412        }
413    }
414}
415
416#[cfg(test)]
417mod tests {
418    use super::*;
419
420    /// The suggested name follows the format as it steps, and stays a suggestion:
421    /// typing still replaces it whole.
422    #[test]
423    fn a_suggested_path_follows_the_format() {
424        let mut modal = ExportModal::new();
425        modal.selected_format = ExportFormat::Csv;
426        modal.suggest_path("people-export");
427        assert_eq!(modal.path_input.value(), "people-export.csv");
428        modal.step_format(1);
429        assert_eq!(
430            modal.path_input.value(),
431            format!("people-export.{}", modal.selected_format.extension())
432        );
433        assert!(modal.path_input.is_suggested());
434    }
435
436    #[test]
437    fn from_path_reads_the_extension_and_looks_through_compression() {
438        assert_eq!(ExportFormat::from_path("out.csv"), Some(ExportFormat::Csv));
439        assert_eq!(
440            ExportFormat::from_path("a/b/out.PARQUET"),
441            Some(ExportFormat::Parquet)
442        );
443        assert_eq!(
444            ExportFormat::from_path("out.csv.gz"),
445            Some(ExportFormat::Csv)
446        );
447        assert_eq!(
448            ExportFormat::from_path("out.jsonl"),
449            Some(ExportFormat::Ndjson)
450        );
451        assert_eq!(ExportFormat::from_path("out"), None);
452        assert_eq!(ExportFormat::from_path("out.dat"), None);
453        // A bare compression extension names no format either way.
454        assert_eq!(ExportFormat::from_path("out.gz"), None);
455    }
456
457    /// TSV and PSV are CSV presets: their names pick them, their delimiter is set, and
458    /// the form offers the header and compression without a delimiter row.
459    #[test]
460    fn tsv_and_psv_are_presets_with_their_delimiter_set() {
461        assert_eq!(ExportFormat::from_path("out.tsv"), Some(ExportFormat::Tsv));
462        assert_eq!(
463            ExportFormat::from_path("out.psv.zst"),
464            Some(ExportFormat::Psv)
465        );
466        assert_eq!(ExportFormat::Tsv.preset_delimiter(), Some(b'\t'));
467        assert_eq!(ExportFormat::Psv.preset_delimiter(), Some(b'|'));
468        assert_eq!(ExportFormat::Csv.preset_delimiter(), None);
469        let mut modal = ExportModal::new();
470        for format in [ExportFormat::Tsv, ExportFormat::Psv] {
471            modal.selected_format = format;
472            let order = modal.focus_order();
473            assert!(!order.contains(&ExportFocus::CsvDelimiter), "{format:?}");
474            assert!(order.contains(&ExportFocus::CsvIncludeHeader), "{format:?}");
475            assert!(order.contains(&ExportFocus::Compression), "{format:?}");
476        }
477        modal.selected_format = ExportFormat::Csv;
478        modal.path_input.set_value("out.csv");
479        modal.selected_format = ExportFormat::Tsv;
480        modal.sync_path_to_format();
481        assert_eq!(modal.path_input.value(), "out.tsv");
482    }
483
484    #[test]
485    fn sync_format_follows_a_known_extension_and_only_that() {
486        let mut modal = ExportModal::new();
487        modal.selected_format = ExportFormat::Parquet;
488        modal.path_input.set_value("out.csv");
489        modal.sync_format_to_path();
490        assert_eq!(modal.selected_format, ExportFormat::Csv);
491
492        modal.selected_format = ExportFormat::Parquet;
493        modal.path_input.set_value("out.dat");
494        modal.sync_format_to_path();
495        assert_eq!(modal.selected_format, ExportFormat::Parquet);
496    }
497
498    #[test]
499    fn a_compression_suffix_in_the_path_sets_compression() {
500        let mut modal = ExportModal::new();
501        modal.path_input.set_value("out.csv.gz");
502        modal.sync_format_to_path();
503        assert_eq!(modal.selected_format, ExportFormat::Csv);
504        assert_eq!(modal.csv_compression, Some(CompressionFormat::Gzip));
505
506        let mut modal = ExportModal::new();
507        modal.path_input.set_value("out.jsonl.zst");
508        modal.sync_format_to_path();
509        assert_eq!(modal.selected_format, ExportFormat::Ndjson);
510        assert_eq!(modal.ndjson_compression, Some(CompressionFormat::Zstd));
511
512        // A plain path leaves an explicitly chosen compression standing.
513        let mut modal = ExportModal::new();
514        modal.csv_compression = Some(CompressionFormat::Gzip);
515        modal.path_input.set_value("out.csv");
516        modal.sync_format_to_path();
517        assert_eq!(modal.csv_compression, Some(CompressionFormat::Gzip));
518    }
519
520    /// A field the new format does not show never keeps focus: stepping the
521    /// format settles it on one that is shown.
522    #[test]
523    fn focus_stays_on_a_shown_field_as_the_format_steps() {
524        let mut modal = ExportModal::new();
525        for format in ExportFormat::ALL {
526            modal.selected_format = format;
527            for field in modal.focus_order() {
528                modal.selected_format = format;
529                crate::form::Form::set_focused(&mut modal, field);
530                for delta in [1, -1, 1, 1] {
531                    modal.step_format(delta);
532                    assert!(
533                        modal.focus_order().contains(&modal.focus),
534                        "{field:?} from {format:?} lands on {:?} at {:?}",
535                        modal.focus,
536                        modal.selected_format
537                    );
538                }
539            }
540        }
541    }
542
543    #[test]
544    fn picking_a_format_rewrites_the_typed_extension() {
545        let mut modal = ExportModal::new();
546        modal.path_input.set_value("out.csv");
547        modal.selected_format = ExportFormat::Parquet;
548        modal.sync_path_to_format();
549        assert_eq!(modal.path_input.value(), "out.parquet");
550
551        // The compression suffix goes when the new format cannot carry it...
552        let mut modal = ExportModal::new();
553        modal.path_input.set_value("data.v2.csv.gz");
554        modal.selected_format = ExportFormat::Parquet;
555        modal.sync_path_to_format();
556        assert_eq!(modal.path_input.value(), "data.v2.parquet");
557
558        // ...and stays when it can.
559        let mut modal = ExportModal::new();
560        modal.path_input.set_value("out.csv.gz");
561        modal.selected_format = ExportFormat::Ndjson;
562        modal.sync_path_to_format();
563        assert_eq!(modal.path_input.value(), "out.jsonl.gz");
564
565        // An extension naming no format is the user's to keep.
566        let mut modal = ExportModal::new();
567        modal.path_input.set_value("out.dat");
568        modal.selected_format = ExportFormat::Parquet;
569        modal.sync_path_to_format();
570        assert_eq!(modal.path_input.value(), "out.dat");
571    }
572}