src/main.rs

#![deny(missing_docs)]
//! Command-line entry point for DADOES.

use std::io::Read;
use std::path::Path;
use std::process::ExitCode;

use dadoes::{DadoesClassifier, EmotionClassifier, ModelIoError, MoodAnalysis};

const DEFAULT_MODEL_PATH: &str = "models/goemotions-linear.dadoes";

/// Runs the DADOES command-line classifier.
fn main() -> ExitCode {
    match run() {
        Ok(()) => ExitCode::SUCCESS,
        Err(error) => {
            eprintln!("{error}");
            ExitCode::from(1)
        }
    }
}

fn run() -> Result<(), CliError> {
    let text = input_text()?;
    if text.trim().is_empty() {
        eprintln!("usage: dadoes <daily report text>");
        return Ok(());
    }

    let classifier = load_classifier()?;
    let analysis = classifier.classify(&text);
    println!("{}", analysis_to_json(&analysis));
    Ok(())
}

fn load_classifier() -> Result<DadoesClassifier, CliError> {
    let model_path = Path::new(DEFAULT_MODEL_PATH);
    if model_path.exists() {
        return DadoesClassifier::from_path(model_path).map_err(CliError::Checkpoint);
    }

    DadoesClassifier::from_default_model().map_err(CliError::Checkpoint)
}

fn input_text() -> Result<String, CliError> {
    let mut arguments = std::env::args().skip(1);
    let first = arguments.next();

    let Some(first) = first else {
        let mut input = String::new();
        std::io::stdin().read_to_string(&mut input)?;
        return Ok(input);
    };

    let mut text = first;
    for argument in arguments {
        text.push(' ');
        text.push_str(&argument);
    }

    Ok(text)
}

fn analysis_to_json(analysis: &MoodAnalysis) -> String {
    let primary = analysis.primary_mood();
    let mut json = String::from("{\"primary_mood\":");

    match primary {
        Some(score) => {
            json.push('"');
            json.push_str(score.mood.as_str());
            json.push('"');
        }
        None => json.push_str("null"),
    }

    json.push_str(",\"moods\":[");

    for (position, score) in analysis.scores().iter().enumerate() {
        if position > 0 {
            json.push(',');
        }

        json.push_str("{\"label\":\"");
        json.push_str(score.mood.as_str());
        json.push_str("\",\"score\":");
        json.push_str(&format!("{:.4}", score.score));
        json.push('}');
    }

    json.push_str("]}");
    json
}

#[derive(Debug)]
enum CliError {
    Io(std::io::Error),
    Checkpoint(ModelIoError),
}

impl std::fmt::Display for CliError {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            CliError::Io(error) => write!(formatter, "failed to read input: {error}"),
            CliError::Checkpoint(error) => write!(formatter, "failed to load checkpoint: {error}"),
        }
    }
}

impl From<std::io::Error> for CliError {
    fn from(error: std::io::Error) -> Self {
        CliError::Io(error)
    }
}