1use 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 Tsv,
13 Psv,
15 Parquet,
16 Json,
17 Ndjson,
18 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 pub fn holds_nesting(self) -> bool {
78 !self.is_delimited()
79 }
80
81 pub fn is_delimited(self) -> bool {
83 matches!(self, Self::Csv | Self::Tsv | Self::Psv)
84 }
85
86 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 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 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#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
129pub enum ExportFocus {
130 #[default]
131 FormatSelector,
132 PathInput,
133 CsvDelimiter,
134 CsvIncludeHeader,
135 Compression,
137 SourceFile,
141}
142
143pub 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 pub csv_delimiter_input: TextInput,
159 pub csv_include_header: bool,
160 pub source_file: bool,
162 pub offer_source_file: bool,
164 pub nested_columns: bool,
167 pub avro_renames: bool,
170 pub csv_compression: Option<CompressionFormat>,
171 pub json_compression: Option<CompressionFormat>,
173 pub ndjson_compression: Option<CompressionFormat>,
175 pub history_limit: usize,
176 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 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 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 pub fn suspend(&mut self) {
233 self.active = false;
234 }
235
236 pub fn resume(&mut self) {
237 self.active = true;
238 }
239
240 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 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 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 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 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 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 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 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 crate::form::Form::settle_focus(self);
342 }
343
344 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 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 #[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 assert_eq!(ExportFormat::from_path("out.gz"), None);
455 }
456
457 #[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 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 #[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 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 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 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}