virittaa

audio metadata analyzer
Log | Files | Refs | README | LICENSE

main.rs (10970B)

      1 pub mod audio_processing;
      2 pub mod constants;
      3 pub mod db;
      4 pub mod file_io;
      5 pub mod import;
      6 pub mod types;
      7 pub mod utils;
      8 
      9 use anyhow::{Context, Result};
     10 use clap::{Arg, Command};
     11 use rayon::prelude::*;
     12 use regex::Regex;
     13 use std::path::PathBuf;
     14 use walkdir::WalkDir;
     15 
     16 use crate::db::LibraryDb;
     17 use crate::db::TrackReport;
     18 use crate::file_io::process_file;
     19 use crate::import::import_mixxx;
     20 use crate::types::TrackError;
     21 use crate::utils::{format_primary_key, is_supported_audio_file, print_entries};
     22 
     23 fn build_cli() -> Command {
     24     Command::new("virittaa")
     25         .version(clap::crate_version!())
     26         .about("Analyze audio files to detect BPM and Key")
     27         .arg(
     28             Arg::new("path")
     29                 .num_args(1..)
     30                 .help("Files or directories to analyze"),
     31         )
     32         .arg(
     33             Arg::new("write-tags")
     34                 .short('t')
     35                 .long("write-tags")
     36                 .action(clap::ArgAction::SetTrue)
     37                 .help("Write detected BPM and Key to file metadata"),
     38         )
     39         .arg(
     40             Arg::new("no-store")
     41                 .long("no-store")
     42                 .action(clap::ArgAction::SetTrue)
     43                 .help("Skip saving track reports to the database"),
     44         )
     45         .arg(
     46             Arg::new("force")
     47                 .short('f')
     48                 .long("force")
     49                 .action(clap::ArgAction::SetTrue)
     50                 .help("Re-analyze files that already have BPM/Key tags"),
     51         )
     52         .arg(
     53             Arg::new("jobs")
     54                 .short('j')
     55                 .long("jobs")
     56                 .help("Number of threads (0 = all cores)")
     57                 .default_value("0")
     58                 .value_parser(clap::value_parser!(usize)),
     59         )
     60         .arg(
     61             Arg::new("list")
     62                 .short('l')
     63                 .long("list")
     64                 .action(clap::ArgAction::SetTrue)
     65                 .help("List track reports in the database"),
     66         )
     67         .arg(
     68             Arg::new("limit")
     69                 .long("limit")
     70                 .help("Max entries to list (default: 1000, 0 = all)")
     71                 .default_value("1000")
     72                 .value_parser(clap::value_parser!(usize)),
     73         )
     74         .arg(
     75             Arg::new("query")
     76                 .long("query")
     77                 .help("Look up a track by exact key"),
     78         )
     79         .arg(
     80             Arg::new("search")
     81                 .short('S')
     82                 .long("search")
     83                 .help("Search track reports by regex"),
     84         )
     85         .arg(
     86             Arg::new("artist")
     87                 .short('a')
     88                 .long("artist")
     89                 .help("Find tracks by artist name prefix"),
     90         )
     91         .arg(
     92             Arg::new("key")
     93                 .long("key")
     94                 .help("Find tracks by key prefix (e.g. 'C Major', 'Am')"),
     95         )
     96         .arg(
     97             Arg::new("bpm")
     98                 .long("bpm")
     99                 .num_args(2)
    100                 .help("Find tracks with BPM in range MIN MAX"),
    101         )
    102         .arg(
    103             Arg::new("import-mixxx")
    104                 .long("import-mixxx")
    105                 .help("Import tracks from a Mixxx SQLite database"),
    106         )
    107         .arg(
    108             Arg::new("db-path")
    109                 .long("db-path")
    110                 .help("Database directory (default: ~/.local/share/virittaa)"),
    111         )
    112 }
    113 
    114 #[allow(clippy::too_many_lines)]
    115 fn main() -> Result<()> {
    116     // Sled's B-tree operations can overflow the default stack on large databases.
    117     std::thread::Builder::new()
    118         .stack_size(32 * 1024 * 1024)
    119         .spawn(real_main)
    120         .expect("Failed to spawn main thread")
    121         .join()
    122         .expect("Main thread panicked")
    123 }
    124 
    125 #[allow(clippy::too_many_lines)]
    126 fn real_main() -> Result<()> {
    127     let mut cmd = build_cli();
    128 
    129     if std::env::args().len() == 1 {
    130         let _ = cmd.print_help();
    131         return Ok(());
    132     }
    133 
    134     let matches = cmd.get_matches();
    135 
    136     let write_tags = matches.get_flag("write-tags");
    137     let no_store = matches.get_flag("no-store");
    138     let force = matches.get_flag("force");
    139     let jobs = *matches.get_one::<usize>("jobs").unwrap();
    140     let list_db = matches.get_flag("list");
    141     let query_key = matches.get_one::<String>("query").map(String::as_str);
    142     let search_pattern = matches.get_one::<String>("search");
    143     let artist_prefix = matches.get_one::<String>("artist").map(String::as_str);
    144     let key_prefix = matches.get_one::<String>("key").map(String::as_str);
    145     let import_mixxx_path = matches.get_one::<String>("import-mixxx");
    146     let limit = *matches.get_one::<usize>("limit").unwrap();
    147     let bpm_range: Option<(f64, f64)> = match matches.get_many::<String>("bpm") {
    148         Some(vals) => {
    149             let v: Vec<&String> = vals.collect();
    150             let min = v[0]
    151                 .parse::<f64>()
    152                 .with_context(|| format!("Invalid BPM min value: {}", v[0]))?;
    153             let max = v[1]
    154                 .parse::<f64>()
    155                 .with_context(|| format!("Invalid BPM max value: {}", v[1]))?;
    156             Some((min, max))
    157         }
    158         None => None,
    159     };
    160 
    161     let db_path = matches.get_one::<String>("db-path").map_or_else(
    162         || {
    163             dirs::data_dir().map_or_else(
    164                 || "~/.local/share/virittaa".to_string(),
    165                 |p| p.join("virittaa").to_string_lossy().to_string(),
    166             )
    167         },
    168         String::from,
    169     );
    170 
    171     let is_db_query = list_db
    172         || query_key.is_some()
    173         || search_pattern.is_some()
    174         || artist_prefix.is_some()
    175         || key_prefix.is_some()
    176         || bpm_range.is_some()
    177         || import_mixxx_path.is_some();
    178 
    179     if is_db_query {
    180         let db = LibraryDb::open(std::path::Path::new(&db_path))?;
    181 
    182         if let Some(mixxx_path) = import_mixxx_path {
    183             let count = import_mixxx(&db, mixxx_path)?;
    184             println!("Imported {count} tracks from Mixxx database.");
    185             return Ok(());
    186         }
    187 
    188         if let Some(query_key) = query_key {
    189             match db.get(query_key)? {
    190                 Some(report) => println!("{query_key}: {report}"),
    191                 None => eprintln!("No track found for key: {query_key}"),
    192             }
    193         } else if let Some(prefix) = artist_prefix {
    194             let entries = db.find_by_artist(prefix, limit)?;
    195             print_entries(
    196                 &entries,
    197                 &format!("No tracks found for artist prefix: {prefix}"),
    198                 "track(s) found.",
    199             );
    200         } else if let Some(prefix) = key_prefix {
    201             let entries = db.find_by_key(prefix, limit)?;
    202             print_entries(
    203                 &entries,
    204                 &format!("No tracks found for key prefix: {prefix}"),
    205                 "track(s) found.",
    206             );
    207         } else if let Some((min, max)) = bpm_range {
    208             let entries = db.find_by_bpm_range(min, max, limit)?;
    209             print_entries(
    210                 &entries,
    211                 &format!("No tracks found with BPM between {min} and {max}"),
    212                 "track(s) found.",
    213             );
    214         } else if let Some(pattern) = search_pattern {
    215             let re = Regex::new(pattern)?;
    216             let entries = db.search(&re, limit)?;
    217             print_entries(
    218                 &entries,
    219                 &format!("No matches for: {pattern}"),
    220                 "match(es) found.",
    221             );
    222         } else {
    223             let entries = db.list(limit)?;
    224             let total = db.count()?;
    225             if entries.is_empty() {
    226                 println!("No track reports found in database.");
    227             } else {
    228                 for (key, report) in &entries {
    229                     println!("{key}: {report}");
    230                 }
    231                 if total > entries.len() {
    232                     println!("\n{} of {} track(s) shown.", entries.len(), total);
    233                 } else {
    234                     println!("\n{} track(s) in database.", entries.len());
    235                 }
    236             }
    237         }
    238         return Ok(());
    239     }
    240 
    241     let paths: Vec<String> = matches
    242         .get_many::<String>("path")
    243         .ok_or_else(|| {
    244             anyhow::anyhow!(
    245                 "Provide file paths, or use --list / --query / --search / --artist / --key / --bpm / --import-mixxx to read the database"
    246             )
    247         })?
    248         .cloned()
    249         .collect();
    250 
    251     if jobs > 0 {
    252         rayon::ThreadPoolBuilder::new()
    253             .num_threads(jobs)
    254             .build_global()
    255             .expect("Failed to build thread pool");
    256     }
    257 
    258     let report_db = if no_store {
    259         None
    260     } else {
    261         Some(LibraryDb::open(std::path::Path::new(&db_path))?)
    262     };
    263 
    264     let files_to_process: Vec<PathBuf> = paths
    265         .into_iter()
    266         .flat_map(|path_str| {
    267             let path = PathBuf::from(&path_str);
    268             if path.is_dir() {
    269                 WalkDir::new(path)
    270                     .into_iter()
    271                     .filter_map(Result::ok)
    272                     .filter(|entry| is_supported_audio_file(entry.path()))
    273                     .map(|entry| entry.path().to_path_buf())
    274                     .collect()
    275             } else if path.is_file() {
    276                 vec![path]
    277             } else {
    278                 vec![]
    279             }
    280         })
    281         .collect::<Vec<_>>();
    282 
    283     let (results, errors): (Vec<_>, Vec<_>) = files_to_process
    284         .par_iter()
    285         .map(|path| {
    286             process_file(path, write_tags, force).map_err(|reason| TrackError {
    287                 path: path.clone(),
    288                 reason,
    289             })
    290         })
    291         .partition(Result::is_ok);
    292 
    293     let results: Vec<_> = results.into_iter().map(|r| r.unwrap()).collect();
    294     let errors: Vec<_> = errors.into_iter().map(|r| r.unwrap_err()).collect();
    295 
    296     if !errors.is_empty() {
    297         for err in &errors {
    298             eprintln!("Error: {} - {}", err.path.display(), err.reason);
    299         }
    300     }
    301 
    302     if let Some(db) = &report_db {
    303         let entries: Vec<(String, TrackReport)> = results
    304             .iter()
    305             .map(|track| {
    306                 let key = format_primary_key(track.artist.as_deref(), track.track.as_deref());
    307                 let report = TrackReport {
    308                     artist: track.artist.clone(),
    309                     track: track.track.clone(),
    310                     bpm: Some(f64::from(track.bpm)),
    311                     key: Some(track.key.clone()),
    312                     timestamp: Some(
    313                         std::time::SystemTime::now()
    314                             .duration_since(std::time::UNIX_EPOCH)
    315                             .unwrap_or_default()
    316                             .as_secs()
    317                             .cast_signed(),
    318                     ),
    319                 };
    320                 (key, report)
    321             })
    322             .collect();
    323 
    324         db.save_batch(&entries)?;
    325         db.flush()?;
    326 
    327         for report in entries.iter().map(|(_, r)| r) {
    328             println!("  Saved: {report}");
    329         }
    330     }
    331 
    332     Ok(())
    333 }