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::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 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,
140}
141
142pub 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 pub source_file: bool,
159 pub offer_source_file: bool,
161 pub nested_columns: bool,
164 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 pub path_error: Option<String>,
174 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 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 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 pub(crate) fn forget_counts(&mut self) {
231 self.counts = None;
232 }
233
234 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 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 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 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 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 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 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 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 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 crate::app::form::Form::settle_focus(self);
333 }
334
335 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 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;