Skip to main content

alimentar/cli/
mod.rs

1//! alimentar CLI - Data Loading, Distribution and Tooling
2//!
3//! Command-line interface for alimentar operations.
4
5use 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
17// Re-export subcommand enums
18pub 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/// alimentar - Data Loading, Distribution and Tooling in Pure Rust
27#[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/// The alimentar command surface.
36///
37/// PUBLIC so `apr data` can expose it. APR-MONO consolidated this crate in-tree,
38/// but the capability stayed reachable only through the standalone `alimentar`
39/// binary: `apr data` shipped 5 commands (audit, split, decontaminate, dedup,
40/// balance) against alimentar's 20, so 18 were unreachable from `apr` at all.
41/// Deleting the binary before exposing these would have removed the capability
42/// rather than relocating it.
43#[derive(Subcommand, Debug)]
44pub enum Commands {
45    /// Convert between data formats
46    Convert {
47        /// Input file path
48        input: PathBuf,
49        /// Output file path
50        output: PathBuf,
51    },
52    /// Display dataset information
53    Info {
54        /// Path to dataset file
55        path: PathBuf,
56    },
57    /// Display first N rows of a dataset
58    Head {
59        /// Path to dataset file
60        path: PathBuf,
61        /// Number of rows to display
62        #[arg(short = 'n', long, default_value = "10")]
63        rows: usize,
64    },
65    /// Display dataset schema
66    Schema {
67        /// Path to dataset file
68        path: PathBuf,
69    },
70    /// Mix multiple datasets with weighted sampling
71    Mix {
72        /// Input files with optional weights (file:weight, e.g.,
73        /// "data.parquet:0.8")
74        #[arg(required = true)]
75        inputs: Vec<String>,
76        /// Output file path
77        #[arg(short, long)]
78        output: PathBuf,
79        /// Random seed for reproducibility
80        #[arg(short, long, default_value = "42")]
81        seed: u64,
82        /// Maximum total rows in output (0 = sum of all weighted inputs)
83        #[arg(short = 'n', long, default_value = "0")]
84        max_rows: usize,
85    },
86    /// Apply Fill-in-the-Middle (FIM) transform for code model training
87    #[cfg(feature = "shuffle")]
88    Fim {
89        /// Input dataset file (Parquet/CSV/JSON)
90        input: PathBuf,
91        /// Output file path
92        #[arg(short, long)]
93        output: PathBuf,
94        /// Column containing code text
95        #[arg(long, default_value = "text")]
96        column: String,
97        /// FIM application rate (0.0-1.0)
98        #[arg(long, default_value = "0.5")]
99        rate: f64,
100        /// FIM format: psm or spm
101        #[arg(long, default_value = "psm")]
102        format: String,
103        /// Random seed for reproducibility
104        #[arg(long, default_value = "42")]
105        seed: u64,
106    },
107    /// Deduplicate dataset by text content (R-019)
108    Dedup {
109        /// Input dataset file
110        input: PathBuf,
111        /// Output file path
112        #[arg(short, long)]
113        output: PathBuf,
114        /// Column to dedup on (auto-detected if not specified)
115        #[arg(long)]
116        column: Option<String>,
117    },
118    /// Filter dataset by text quality signals (R-022)
119    #[command(name = "filter-text")]
120    FilterText {
121        /// Input dataset file
122        input: PathBuf,
123        /// Output file path
124        #[arg(short, long)]
125        output: PathBuf,
126        /// Column containing text (auto-detected if not specified)
127        #[arg(long)]
128        column: Option<String>,
129        /// Minimum composite quality score (0.0-1.0)
130        #[arg(long, default_value = "0.4")]
131        min_score: f64,
132        /// Minimum document length in characters
133        #[arg(long, default_value = "50")]
134        min_length: usize,
135        /// Maximum document length in characters
136        #[arg(long, default_value = "1000000")]
137        max_length: usize,
138    },
139    /// Interactive TUI viewer for datasets
140    View {
141        /// Path to dataset file (Parquet/Arrow/CSV/JSON)
142        path: PathBuf,
143        /// Initial search query
144        #[arg(long)]
145        search: Option<String>,
146    },
147    /// Import dataset from local files or HuggingFace Hub
148    Import {
149        #[command(subcommand)]
150        source: ImportSource,
151    },
152    /// HuggingFace Hub commands (push/upload datasets)
153    #[allow(clippy::doc_markdown)]
154    #[cfg(feature = "hf-hub")]
155    #[command(subcommand)]
156    Hub(HubCommands),
157    /// Registry commands for dataset sharing and discovery
158    #[command(subcommand)]
159    Registry(RegistryCommands),
160    /// Data drift detection commands
161    #[command(subcommand)]
162    Drift(DriftCommands),
163    /// Data quality checking commands
164    #[command(subcommand)]
165    Quality(QualityCommands),
166    /// Federated split coordination commands
167    #[command(subcommand)]
168    Fed(FedCommands),
169    /// Python doctest extraction commands
170    #[cfg(feature = "doctest")]
171    #[command(subcommand)]
172    Doctest(DoctestCommands),
173    /// Interactive REPL for data exploration
174    #[cfg(feature = "repl")]
175    Repl,
176}
177
178/// Python doctest extraction commands
179#[cfg(feature = "doctest")]
180#[derive(Subcommand, Debug)]
181pub enum DoctestCommands {
182    /// Extract doctests from Python source files
183    ///
184    /// apr propagates its auto `--version` into every subcommand
185    /// (`propagate_version = true`), which collides with this real
186    /// `--version` argument, so the auto flag is disabled here.
187    #[command(disable_version_flag = true)]
188    Extract {
189        /// Input directory containing Python source files
190        input: PathBuf,
191        /// Output parquet file
192        #[arg(short, long)]
193        output: PathBuf,
194        /// Source identifier (e.g., "cpython", "numpy")
195        #[arg(short, long, default_value = "unknown")]
196        source: String,
197        /// Version string or git SHA
198        // Long-only: apr propagates its global -v/--verbose and -q/--quiet
199        // into every subcommand, so a derived short here makes the whole
200        // clap tree invalid and the subcommand panics on any invocation.
201        #[arg(long, default_value = "unknown")]
202        version: String,
203    },
204    /// Merge multiple doctest corpora into one
205    Merge {
206        /// Input parquet files to merge
207        #[arg(required = true)]
208        inputs: Vec<PathBuf>,
209        /// Output parquet file
210        #[arg(short, long)]
211        output: PathBuf,
212    },
213}
214
215#[allow(clippy::too_many_lines)]
216/// Run the alimentar CLI.
217pub fn run() -> ExitCode {
218    dispatch(Cli::parse().command)
219}
220
221/// Execute one alimentar command.
222///
223/// Split out of `run()` so `apr data` dispatches the SAME implementation
224/// rather than re-declaring the clap tree or shelling out to a second binary
225/// -- one implementation, two names.
226#[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            &registry,
349        ),
350        RegistryCommands::Pull {
351            name,
352            output,
353            version,
354            registry,
355        } => registry::cmd_registry_pull(&name, &output, version.as_deref(), &registry),
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, &current, &tests, alpha, &format),
375        DriftCommands::Report {
376            reference,
377            current,
378            output,
379        } => drift::cmd_drift_report(&reference, &current, 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, &current, 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// =============================================================================
477// Doctest Commands
478// =============================================================================
479
480#[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    // Load all datasets and concatenate
527    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    // Create merged dataset
543    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}