janitor

janitor
git clone https://git.prime8.dev/janitor.git
Log | Files | Refs | README | LICENSE

main.rs (4274B)

      1 use anyhow::{Context, Error, Result};
      2 use async_walkdir::{Filtering, WalkDir};
      3 use clap::{Parser, Subcommand};
      4 use spinoff::{spinners, Spinner};
      5 use std::path::{Path, PathBuf};
      6 use tokio::{
      7     fs::{copy, rename},
      8     task::JoinSet,
      9 };
     10 use tokio_stream::StreamExt;
     11 
     12 pub mod audio;
     13 use audio::is_audio_file;
     14 
     15 pub mod processing;
     16 use processing::{get_result_path, process, Label, ResultPathOptions};
     17 
     18 #[derive(Parser)]
     19 #[command(about, long_about = None, version)]
     20 struct Args {
     21     #[command(subcommand)]
     22     command: Option<Command>,
     23 
     24     #[arg(required = true)]
     25     path: PathBuf,
     26 
     27     #[arg(short, long, default_value = "0.0.0.0:8000")]
     28     address: String,
     29 }
     30 
     31 #[derive(Subcommand, Debug)]
     32 pub enum Command {
     33     #[command()]
     34     Copy {
     35         #[arg(short, long)]
     36         speech_dir: Option<PathBuf>,
     37 
     38         #[arg(short, long)]
     39         music_dir: Option<PathBuf>,
     40 
     41         #[arg(short, long)]
     42         noise_dir: Option<PathBuf>,
     43     },
     44 
     45     #[command()]
     46     Move {
     47         #[arg(short, long)]
     48         speech_dir: Option<PathBuf>,
     49 
     50         #[arg(short, long)]
     51         music_dir: Option<PathBuf>,
     52 
     53         #[arg(short, long)]
     54         noise_dir: Option<PathBuf>,
     55     },
     56 }
     57 
     58 impl Command {
     59     async fn perform(&self, path: &Path, label: &Label, options: &ResultPathOptions) -> Result<()> {
     60         if let Some(dir) = get_result_path(path, label, options) {
     61             match self {
     62                 Command::Copy { .. } => {
     63                     copy(path, &dir).await?;
     64                 }
     65                 Command::Move { .. } => {
     66                     rename(path, &dir).await?;
     67                 }
     68             }
     69         }
     70         Ok(())
     71     }
     72 }
     73 
     74 #[tokio::main]
     75 async fn main() -> Result<(), Error> {
     76     let args = Args::parse();
     77 
     78     let url = format!("http://{}/", args.address);
     79 
     80     let options = match args.command {
     81         Some(Command::Copy {
     82             ref speech_dir,
     83             ref music_dir,
     84             ref noise_dir,
     85         })
     86         | Some(Command::Move {
     87             ref speech_dir,
     88             ref music_dir,
     89             ref noise_dir,
     90         }) => Some(ResultPathOptions {
     91             speech_dir: speech_dir.clone(),
     92             music_dir: music_dir.clone(),
     93             noise_dir: noise_dir.clone(),
     94         }),
     95         _ => None,
     96     };
     97 
     98     let mut spinner = Spinner::new(spinners::Line, "Loading...", None);
     99 
    100     if args.path.is_file() {
    101         spinner.update_text(format!("Labelling {:?}", args.path));
    102         let (path, label) = process(args.path.clone(), url.clone())
    103             .await
    104             .with_context(|| format!("Failed to process {:?}", args.path))?;
    105 
    106         if let Some(ref command) = args.command {
    107             let opts = options.as_ref().context("Missing result path options")?;
    108             command
    109                 .perform(&path, &label, opts)
    110                 .await
    111                 .with_context(|| format!("Failed to perform command {:?}", command))?;
    112         } else {
    113             spinner.stop_with_message(&format!("{:?}: {:?}", path, label));
    114         }
    115 
    116         return Ok(());
    117     }
    118 
    119     let mut entries = WalkDir::new(args.path).filter(|entry| async move {
    120         if is_audio_file(entry).await {
    121             Filtering::Continue
    122         } else {
    123             Filtering::Ignore
    124         }
    125     });
    126 
    127     let mut jobs = JoinSet::new();
    128 
    129     loop {
    130         match entries.next().await {
    131             Some(Ok(entry)) => {
    132                 let path = entry.path();
    133                 spinner.update_text(format!("Labelling {:?}", path));
    134                 let future = process(path, url.clone());
    135                 jobs.spawn(future);
    136             }
    137             Some(Err(e)) => return Err(e.into()),
    138             None => break,
    139         };
    140     }
    141 
    142     spinner.stop();
    143 
    144     while let Some(result) = jobs.join_next().await {
    145         let (path, label) = result
    146             .context("Task panicked or was cancelled")?
    147             .context("Failed to process file")?;
    148         println!("{:?}: {:?}", path, label);
    149 
    150         if let Some(ref command) = args.command {
    151             let opts = options.as_ref().context("Missing result path options")?;
    152             command
    153                 .perform(&path, &label, opts)
    154                 .await
    155                 .with_context(|| format!("Failed to perform command {:?}", command))?;
    156         }
    157     }
    158     Ok(())
    159 }