use std::fs::File;
use std::io::{BufReader, Write};
use std::process::ExitCode;
use std::sync::Arc;
use serde::Serialize;
use yomitoki::{AnalysisConfig, FragmentCorpus, SynthesizabilityReport, analyze, analyze_smiles};
#[derive(Clone, Copy, PartialEq)]
enum Format {
Human,
Json,
Jsonl,
}
struct Args {
smiles: Option<String>,
input: Option<String>,
output: Option<String>,
format: Format,
fragment_corpus: Option<String>,
}
fn parse_args(mut args: impl Iterator<Item = String>) -> Result<Args, String> {
let mut smiles = None;
let mut input = None;
let mut output = None;
let mut format = None;
let mut fragment_corpus = None;
while let Some(arg) = args.next() {
match arg.as_str() {
"--input" => input = Some(args.next().ok_or("--input requires a value")?),
"--output" => output = Some(args.next().ok_or("--output requires a value")?),
"--fragment-corpus" => {
fragment_corpus = Some(args.next().ok_or("--fragment-corpus requires a value")?)
}
"--format" => {
let value = args.next().ok_or("--format requires a value")?;
format = Some(match value.as_str() {
"human" => Format::Human,
"json" => Format::Json,
"jsonl" => Format::Jsonl,
other => {
return Err(format!(
"unknown format {other:?} (expected human|json|jsonl)"
));
}
});
}
other if !other.starts_with("--") && smiles.is_none() && input.is_none() => {
smiles = Some(other.to_string());
}
other => return Err(format!("unrecognized argument {other:?}")),
}
}
if smiles.is_some() && input.is_some() {
return Err("provide either a SMILES argument or --input, not both".to_string());
}
if smiles.is_none() && input.is_none() {
return Err("provide a SMILES argument or --input <file>".to_string());
}
Ok(Args {
smiles,
input,
output,
format: format.unwrap_or(Format::Human),
fragment_corpus,
})
}
const USAGE: &str = "\
Usage:
yomitoki analyze \"<SMILES>\" [--format human|json|jsonl] [--fragment-corpus <dir>]
yomitoki analyze --input <file> [--format human|json|jsonl] [--output <file>] [--fragment-corpus <dir>]
<file> may be a .sdf file or a SMILES-per-line file (optionally with a
whitespace-separated name column).
<dir> for --fragment-corpus is a tools/build-fragment-corpus output
directory (containing fragment_frequencies.json and manifest.json).
Enables the fragment_precedent component; omitted, it stays None.";
#[derive(Serialize)]
struct BatchItem<'a> {
input: &'a str,
#[serde(skip_serializing_if = "Option::is_none")]
report: Option<SynthesizabilityReport>,
#[serde(skip_serializing_if = "Option::is_none")]
error: Option<String>,
}
fn main() -> ExitCode {
let mut raw_args = std::env::args().skip(1);
match raw_args.next().as_deref() {
Some("analyze") => {}
Some("--help") | Some("-h") | None => {
println!("{USAGE}");
return ExitCode::SUCCESS;
}
Some(other) => {
eprintln!("unknown subcommand {other:?}\n\n{USAGE}");
return ExitCode::from(2);
}
}
let args = match parse_args(raw_args) {
Ok(args) => args,
Err(message) => {
eprintln!("error: {message}\n\n{USAGE}");
return ExitCode::from(2);
}
};
let mut writer: Box<dyn Write> = match &args.output {
Some(path) => match File::create(path) {
Ok(file) => Box::new(file),
Err(e) => {
eprintln!("error: could not create output file {path:?}: {e}");
return ExitCode::FAILURE;
}
},
None => Box::new(std::io::stdout()),
};
let mut config = AnalysisConfig::default();
if let Some(dir) = &args.fragment_corpus {
match FragmentCorpus::load_dir(dir) {
Ok(corpus) => config.fragment_model.corpus = Some(Arc::new(corpus)),
Err(e) => {
eprintln!("error: could not load --fragment-corpus {dir:?}: {e}");
return ExitCode::FAILURE;
}
}
}
if let Some(smiles) = &args.smiles {
return run_single(smiles, &config, args.format, writer.as_mut());
}
let path = args.input.as_deref().expect("checked in parse_args");
run_batch(path, &config, args.format, writer.as_mut())
}
fn run_single(
smiles: &str,
config: &AnalysisConfig,
format: Format,
writer: &mut dyn Write,
) -> ExitCode {
let report = match analyze_smiles(smiles, config) {
Ok(report) => report,
Err(e) => {
eprintln!("error: {e}");
return ExitCode::FAILURE;
}
};
let result = match format {
Format::Human => writeln!(writer, "{}", render_human(&report)),
Format::Json => writeln!(
writer,
"{}",
serde_json::to_string_pretty(&report).expect("report always serializes")
),
Format::Jsonl => {
let item = BatchItem {
input: smiles,
report: Some(report),
error: None,
};
writeln!(
writer,
"{}",
serde_json::to_string(&item).expect("item always serializes")
)
}
};
if let Err(e) = result {
eprintln!("error: failed to write output: {e}");
return ExitCode::FAILURE;
}
ExitCode::SUCCESS
}
type LabeledMolecule = (String, Result<chematic::core::Molecule, String>);
fn read_batch_records(path: &str) -> Result<Vec<LabeledMolecule>, String> {
if path.to_ascii_lowercase().ends_with(".sdf") {
let content =
std::fs::read_to_string(path).map_err(|e| format!("could not read {path:?}: {e}"))?;
Ok(chematic::mol::SdfReader::new(&content)
.enumerate()
.map(|(i, result)| match result {
Ok((mol, meta)) => {
let label = if meta.name.is_empty() {
format!("record {}", i + 1)
} else {
meta.name
};
(label, Ok(mol))
}
Err(e) => (format!("record {}", i + 1), Err(e.to_string())),
})
.collect())
} else {
let file = File::open(path).map_err(|e| format!("could not open {path:?}: {e}"))?;
let reader = chematic::mol::SmilesRecordReader::new(
BufReader::new(file),
chematic::mol::SmilesReaderOptions::default(),
);
Ok(reader
.enumerate()
.map(|(i, result)| match result {
Ok(record) => {
let label = if record.name.is_empty() {
format!("record {}", i + 1)
} else {
record.name
};
(label, Ok(record.mol))
}
Err(e) => (format!("record {}", i + 1), Err(e.to_string())),
})
.collect())
}
}
fn run_batch(
path: &str,
config: &AnalysisConfig,
format: Format,
writer: &mut dyn Write,
) -> ExitCode {
let records = match read_batch_records(path) {
Ok(records) => records,
Err(message) => {
eprintln!("error: {message}");
return ExitCode::FAILURE;
}
};
let mut any_failed = false;
let mut json_items: Vec<BatchItem> = Vec::with_capacity(records.len());
for (label, molecule_result) in &records {
let outcome = molecule_result
.as_ref()
.map_err(|e| e.clone())
.and_then(|mol| analyze(mol, config).map_err(|e| e.to_string()));
match &outcome {
Ok(_) => {}
Err(_) => any_failed = true,
}
match format {
Format::Human => {
let block = match &outcome {
Ok(report) => render_human(report),
Err(message) => format!("ERROR: {message}"),
};
if let Err(e) = writeln!(writer, "=== {label} ===\n{block}\n") {
eprintln!("error: failed to write output: {e}");
return ExitCode::FAILURE;
}
}
Format::Jsonl => {
let item = BatchItem {
input: label,
report: outcome.as_ref().ok().cloned(),
error: outcome.as_ref().err().cloned(),
};
let line = serde_json::to_string(&item).expect("item always serializes");
if let Err(e) = writeln!(writer, "{line}") {
eprintln!("error: failed to write output: {e}");
return ExitCode::FAILURE;
}
}
Format::Json => {
json_items.push(BatchItem {
input: label,
report: outcome.as_ref().ok().cloned(),
error: outcome.as_ref().err().cloned(),
});
}
}
}
if format == Format::Json {
let rendered = serde_json::to_string_pretty(&json_items).expect("items always serialize");
if let Err(e) = writeln!(writer, "{rendered}") {
eprintln!("error: failed to write output: {e}");
return ExitCode::FAILURE;
}
}
if any_failed {
ExitCode::FAILURE
} else {
ExitCode::SUCCESS
}
}
fn render_human(report: &SynthesizabilityReport) -> String {
let mut out = format!(
"Verdict: {:?}\nSynthesizability: {:.2}\nConfidence: {:.2}",
report.overall.verdict,
report.overall.synthesizability.value(),
report.overall.confidence.value(),
);
if !report.dominant_penalties().is_empty() {
out.push_str("\n\nDominant penalties:");
for (i, penalty) in report.dominant_penalties().iter().enumerate() {
out.push_str(&format!("\n{}. {}", i + 1, penalty.name));
}
}
if !report.suggestions.is_empty() {
out.push_str("\n\nSimplification suggestions (heuristic, not a guarantee):");
for (i, suggestion) in report.suggestions.iter().enumerate() {
out.push_str(&format!(
"\n{}. {:?}: {}",
i + 1,
suggestion.code,
suggestion.rationale
));
}
}
out
}
#[cfg(test)]
mod tests {
use super::*;
fn args(strs: &[&str]) -> std::vec::IntoIter<String> {
strs.iter()
.map(|s| s.to_string())
.collect::<Vec<_>>()
.into_iter()
}
#[test]
fn single_smiles_defaults_to_human_format() {
let parsed = parse_args(args(&["CCO"])).expect("valid");
assert_eq!(parsed.smiles.as_deref(), Some("CCO"));
assert!(parsed.input.is_none());
assert!(parsed.format == Format::Human);
}
#[test]
fn format_flag_is_parsed() {
let parsed = parse_args(args(&["CCO", "--format", "json"])).expect("valid");
assert!(parsed.format == Format::Json);
}
#[test]
fn unknown_format_is_rejected() {
assert!(parse_args(args(&["CCO", "--format", "yaml"])).is_err());
}
#[test]
fn input_and_smiles_together_is_rejected() {
assert!(parse_args(args(&["CCO", "--input", "molecules.smi"])).is_err());
}
#[test]
fn neither_smiles_nor_input_is_rejected() {
assert!(parse_args(args(&["--format", "json"])).is_err());
}
#[test]
fn missing_flag_value_is_rejected() {
assert!(parse_args(args(&["CCO", "--format"])).is_err());
assert!(parse_args(args(&["--input"])).is_err());
assert!(parse_args(args(&["--output"])).is_err());
assert!(parse_args(args(&["CCO", "--fragment-corpus"])).is_err());
}
#[test]
fn fragment_corpus_flag_is_parsed_and_defaults_to_none() {
let without = parse_args(args(&["CCO"])).expect("valid");
assert!(without.fragment_corpus.is_none());
let with = parse_args(args(&["CCO", "--fragment-corpus", "corpus_dir"])).expect("valid");
assert_eq!(with.fragment_corpus.as_deref(), Some("corpus_dir"));
}
#[test]
fn input_output_and_format_are_parsed_together() {
let parsed = parse_args(args(&[
"--input",
"molecules.smi",
"--format",
"jsonl",
"--output",
"reports.jsonl",
]))
.expect("valid");
assert_eq!(parsed.input.as_deref(), Some("molecules.smi"));
assert_eq!(parsed.output.as_deref(), Some("reports.jsonl"));
assert!(parsed.format == Format::Jsonl);
}
#[test]
fn render_human_matches_the_spec_shape() {
let config = AnalysisConfig::default();
let report = analyze_smiles("C1CC2CCC1C2", &config).expect("norbornane");
let text = render_human(&report);
assert!(text.starts_with("Verdict: "));
assert!(text.contains("Synthesizability: "));
assert!(text.contains("Confidence: "));
assert!(text.contains("Dominant penalties:"));
assert!(text.contains("Simplification suggestions"));
}
#[test]
fn render_human_omits_dominant_penalties_and_suggestions_when_empty() {
let config = AnalysisConfig::default();
let report = analyze_smiles("CCO", &config).expect("ethanol");
let text = render_human(&report);
assert!(!text.contains("Dominant penalties"));
assert!(!text.contains("Simplification suggestions"));
}
#[test]
fn single_mode_jsonl_uses_the_same_wrapper_shape_as_batch_mode() {
let config = AnalysisConfig::default();
let mut buffer = Vec::new();
let code = run_single("CCO", &config, Format::Jsonl, &mut buffer);
assert!(code == ExitCode::SUCCESS);
let line = String::from_utf8(buffer).expect("utf8");
let value: serde_json::Value = serde_json::from_str(line.trim()).expect("valid json");
assert_eq!(value["input"], "CCO");
assert!(value.get("report").is_some());
assert!(value.get("error").is_none());
}
}