1use std::{path::PathBuf, process::ExitCode};
6
7use clap::{Parser, Subcommand};
8
9mod basic;
10mod drift;
11mod fed;
12mod hub;
13mod quality;
14mod registry;
15mod view;
16
17pub use drift::DriftCommands;
19pub use fed::FedCommands;
20#[cfg(feature = "hf-hub")]
21pub use hub::HubCommands;
22pub use hub::ImportSource;
23pub use quality::QualityCommands;
24pub use registry::RegistryCommands;
25
26#[derive(Parser)]
28#[command(name = "alimentar")]
29#[command(author, version, about, long_about = None)]
30pub struct Cli {
31 #[command(subcommand)]
32 pub command: Commands,
33}
34
35#[derive(Subcommand, Debug)]
44pub enum Commands {
45 Convert {
47 input: PathBuf,
49 output: PathBuf,
51 },
52 Info {
54 path: PathBuf,
56 },
57 Head {
59 path: PathBuf,
61 #[arg(short = 'n', long, default_value = "10")]
63 rows: usize,
64 },
65 Schema {
67 path: PathBuf,
69 },
70 Mix {
72 #[arg(required = true)]
75 inputs: Vec<String>,
76 #[arg(short, long)]
78 output: PathBuf,
79 #[arg(short, long, default_value = "42")]
81 seed: u64,
82 #[arg(short = 'n', long, default_value = "0")]
84 max_rows: usize,
85 },
86 #[cfg(feature = "shuffle")]
88 Fim {
89 input: PathBuf,
91 #[arg(short, long)]
93 output: PathBuf,
94 #[arg(long, default_value = "text")]
96 column: String,
97 #[arg(long, default_value = "0.5")]
99 rate: f64,
100 #[arg(long, default_value = "psm")]
102 format: String,
103 #[arg(long, default_value = "42")]
105 seed: u64,
106 },
107 Dedup {
109 input: PathBuf,
111 #[arg(short, long)]
113 output: PathBuf,
114 #[arg(long)]
116 column: Option<String>,
117 },
118 #[command(name = "filter-text")]
120 FilterText {
121 input: PathBuf,
123 #[arg(short, long)]
125 output: PathBuf,
126 #[arg(long)]
128 column: Option<String>,
129 #[arg(long, default_value = "0.4")]
131 min_score: f64,
132 #[arg(long, default_value = "50")]
134 min_length: usize,
135 #[arg(long, default_value = "1000000")]
137 max_length: usize,
138 },
139 View {
141 path: PathBuf,
143 #[arg(long)]
145 search: Option<String>,
146 },
147 Import {
149 #[command(subcommand)]
150 source: ImportSource,
151 },
152 #[allow(clippy::doc_markdown)]
154 #[cfg(feature = "hf-hub")]
155 #[command(subcommand)]
156 Hub(HubCommands),
157 #[command(subcommand)]
159 Registry(RegistryCommands),
160 #[command(subcommand)]
162 Drift(DriftCommands),
163 #[command(subcommand)]
165 Quality(QualityCommands),
166 #[command(subcommand)]
168 Fed(FedCommands),
169 #[cfg(feature = "doctest")]
171 #[command(subcommand)]
172 Doctest(DoctestCommands),
173 #[cfg(feature = "repl")]
175 Repl,
176}
177
178#[cfg(feature = "doctest")]
180#[derive(Subcommand, Debug)]
181pub enum DoctestCommands {
182 #[command(disable_version_flag = true)]
188 Extract {
189 input: PathBuf,
191 #[arg(short, long)]
193 output: PathBuf,
194 #[arg(short, long, default_value = "unknown")]
196 source: String,
197 #[arg(long, default_value = "unknown")]
202 version: String,
203 },
204 Merge {
206 #[arg(required = true)]
208 inputs: Vec<PathBuf>,
209 #[arg(short, long)]
211 output: PathBuf,
212 },
213}
214
215#[allow(clippy::too_many_lines)]
216pub fn run() -> ExitCode {
218 dispatch(Cli::parse().command)
219}
220
221#[allow(clippy::too_many_lines)]
227pub fn dispatch(command: Commands) -> ExitCode {
228 let result = match command {
229 Commands::Convert { input, output } => basic::cmd_convert(&input, &output),
230 Commands::Info { path } => basic::cmd_info(&path),
231 Commands::Head { path, rows } => basic::cmd_head(&path, rows),
232 Commands::Schema { path } => basic::cmd_schema(&path),
233 Commands::Mix {
234 inputs,
235 output,
236 seed,
237 max_rows,
238 } => basic::cmd_mix(&inputs, &output, seed, max_rows),
239 #[cfg(feature = "shuffle")]
240 Commands::Fim {
241 input,
242 output,
243 column,
244 rate,
245 format,
246 seed,
247 } => basic::cmd_fim(&input, &output, &column, rate, &format, seed),
248 Commands::Dedup {
249 input,
250 output,
251 column,
252 } => basic::cmd_dedup(&input, &output, column.as_deref()),
253 Commands::FilterText {
254 input,
255 output,
256 column,
257 min_score,
258 min_length,
259 max_length,
260 } => basic::cmd_filter_text(
261 &input,
262 &output,
263 column.as_deref(),
264 min_score,
265 min_length,
266 max_length,
267 ),
268 Commands::View { path, search } => view::cmd_view(&path, search.as_deref()),
269 Commands::Import { source } => match source {
270 ImportSource::Local {
271 input,
272 output,
273 format,
274 } => hub::cmd_import_local(&input, &output, format.as_deref()),
275 #[cfg(feature = "hf-hub")]
276 ImportSource::Hf {
277 repo_id,
278 output,
279 revision,
280 subset,
281 split,
282 } => hub::cmd_import_hf(&repo_id, &output, &revision, subset.as_deref(), &split),
283 },
284 #[cfg(feature = "hf-hub")]
285 Commands::Hub(hub_cmd) => match hub_cmd {
286 HubCommands::Push {
287 input,
288 repo_id,
289 path_in_repo,
290 message,
291 readme,
292 private,
293 } => hub::cmd_hub_push(
294 &input,
295 &repo_id,
296 path_in_repo.as_deref(),
297 &message,
298 readme.as_ref(),
299 private,
300 ),
301 },
302 Commands::Registry(registry_cmd) => dispatch_registry(registry_cmd),
303 Commands::Drift(drift_cmd) => dispatch_drift(drift_cmd),
304 Commands::Quality(quality_cmd) => dispatch_quality(quality_cmd),
305 Commands::Fed(fed_cmd) => dispatch_fed(fed_cmd),
306 #[cfg(feature = "doctest")]
307 Commands::Doctest(doctest_cmd) => match doctest_cmd {
308 DoctestCommands::Extract {
309 input,
310 output,
311 source,
312 version,
313 } => cmd_doctest_extract(&input, &output, &source, &version),
314 DoctestCommands::Merge { inputs, output } => cmd_doctest_merge(&inputs, &output),
315 },
316 #[cfg(feature = "repl")]
317 Commands::Repl => crate::repl::run(),
318 };
319
320 match result {
321 Ok(()) => ExitCode::SUCCESS,
322 Err(e) => {
323 eprintln!("Error: {}", e);
324 ExitCode::FAILURE
325 }
326 }
327}
328
329fn dispatch_registry(cmd: RegistryCommands) -> crate::error::Result<()> {
330 match cmd {
331 RegistryCommands::Init { path } => registry::cmd_registry_init(&path),
332 RegistryCommands::List { path } => registry::cmd_registry_list(&path),
333 RegistryCommands::Push {
334 input,
335 name,
336 version,
337 description,
338 license,
339 tags,
340 registry,
341 } => registry::cmd_registry_push(
342 &input,
343 &name,
344 &version,
345 &description,
346 &license,
347 &tags,
348 ®istry,
349 ),
350 RegistryCommands::Pull {
351 name,
352 output,
353 version,
354 registry,
355 } => registry::cmd_registry_pull(&name, &output, version.as_deref(), ®istry),
356 RegistryCommands::Search { query, path } => registry::cmd_registry_search(&query, &path),
357 RegistryCommands::ShowInfo { name, path } => registry::cmd_registry_show_info(&name, &path),
358 RegistryCommands::Delete {
359 name,
360 version,
361 path,
362 } => registry::cmd_registry_delete(&name, &version, &path),
363 }
364}
365
366fn dispatch_drift(cmd: DriftCommands) -> crate::error::Result<()> {
367 match cmd {
368 DriftCommands::Detect {
369 reference,
370 current,
371 tests,
372 alpha,
373 format,
374 } => drift::cmd_drift_detect(&reference, ¤t, &tests, alpha, &format),
375 DriftCommands::Report {
376 reference,
377 current,
378 output,
379 } => drift::cmd_drift_report(&reference, ¤t, output.as_ref()),
380 DriftCommands::Sketch {
381 input,
382 output,
383 sketch_type,
384 source,
385 format,
386 } => drift::cmd_drift_sketch(&input, &output, &sketch_type, source.as_deref(), &format),
387 DriftCommands::Merge {
388 sketches,
389 output,
390 format,
391 } => drift::cmd_drift_merge(&sketches, &output, &format),
392 DriftCommands::Compare {
393 reference,
394 current,
395 threshold,
396 format,
397 } => drift::cmd_drift_compare(&reference, ¤t, threshold, &format),
398 }
399}
400
401fn dispatch_quality(cmd: QualityCommands) -> crate::error::Result<()> {
402 match cmd {
403 QualityCommands::Check {
404 path,
405 null_threshold,
406 duplicate_threshold,
407 detect_outliers,
408 format,
409 } => quality::cmd_quality_check(
410 &path,
411 null_threshold,
412 duplicate_threshold,
413 detect_outliers,
414 &format,
415 ),
416 QualityCommands::Report { path, output } => {
417 quality::cmd_quality_report(&path, output.as_deref())
418 }
419 QualityCommands::Score {
420 path,
421 profile,
422 suggest,
423 json,
424 badge,
425 } => quality::cmd_quality_score(&path, &profile, suggest, json, badge),
426 QualityCommands::Profiles => quality::cmd_quality_profiles(),
427 }
428}
429
430fn dispatch_fed(cmd: FedCommands) -> crate::error::Result<()> {
431 match cmd {
432 FedCommands::Manifest {
433 input,
434 output,
435 node_id,
436 train_ratio,
437 seed,
438 format,
439 } => fed::cmd_fed_manifest(&input, &output, &node_id, train_ratio, seed, &format),
440 FedCommands::Plan {
441 manifests,
442 output,
443 strategy,
444 train_ratio,
445 seed,
446 stratify_column,
447 format,
448 } => fed::cmd_fed_plan(
449 &manifests,
450 &output,
451 &strategy,
452 train_ratio,
453 seed,
454 stratify_column.as_deref(),
455 &format,
456 ),
457 FedCommands::Split {
458 input,
459 plan,
460 node_id,
461 train_output,
462 test_output,
463 validation_output,
464 } => fed::cmd_fed_split(
465 &input,
466 &plan,
467 &node_id,
468 &train_output,
469 &test_output,
470 validation_output.as_ref(),
471 ),
472 FedCommands::Verify { manifests, format } => fed::cmd_fed_verify(&manifests, &format),
473 }
474}
475
476#[cfg(feature = "doctest")]
481fn cmd_doctest_extract(
482 input: &std::path::Path,
483 output: &std::path::Path,
484 source: &str,
485 version: &str,
486) -> crate::Result<()> {
487 use crate::DocTestParser;
488
489 if !input.is_dir() {
490 return Err(crate::Error::invalid_config(format!(
491 "Input path must be a directory: {}",
492 input.display()
493 )));
494 }
495
496 let parser = DocTestParser::new();
497 let corpus = parser.parse_directory(input, source, version)?;
498
499 println!(
500 "Extracted {} doctests from {} ({})",
501 corpus.len(),
502 source,
503 version
504 );
505
506 if corpus.is_empty() {
507 println!("Warning: No doctests found in {}", input.display());
508 return Ok(());
509 }
510
511 let dataset = corpus.to_dataset()?;
512 dataset.to_parquet(output)?;
513
514 println!("Wrote {} to {}", corpus.len(), output.display());
515 Ok(())
516}
517
518#[cfg(feature = "doctest")]
519fn cmd_doctest_merge(inputs: &[PathBuf], output: &std::path::Path) -> crate::Result<()> {
520 use crate::{dataset::Dataset, ArrowDataset};
521
522 if inputs.is_empty() {
523 return Err(crate::Error::invalid_config("No input files provided"));
524 }
525
526 let mut all_batches = Vec::new();
528 let mut total_rows = 0;
529
530 for input in inputs {
531 let dataset = ArrowDataset::from_parquet(input)?;
532 total_rows += dataset.len();
533 for batch in dataset.iter() {
534 all_batches.push(batch.clone());
535 }
536 }
537
538 if all_batches.is_empty() {
539 return Err(crate::Error::invalid_config("No data found in input files"));
540 }
541
542 let merged = ArrowDataset::new(all_batches)?;
544 merged.to_parquet(output)?;
545
546 println!(
547 "Merged {} doctests from {} files to {}",
548 total_rows,
549 inputs.len(),
550 output.display()
551 );
552 Ok(())
553}