molx 0.1.0

Interactive protein structure exploration in the terminal
Documentation
use std::fs;
use std::path::Path;
use std::time::Duration;

const API_ROOT: &str = "https://api.esmatlas.com";

pub struct AtlasClient {
    client: reqwest::blocking::Client,
}

impl AtlasClient {
    pub fn new() -> Result<Self, String> {
        let client = reqwest::blocking::Client::builder()
            .user_agent(concat!("molx/", env!("CARGO_PKG_VERSION")))
            .connect_timeout(Duration::from_secs(20))
            .timeout(Duration::from_secs(300))
            .build()
            .map_err(|error| format!("cannot initialize HTTP client: {error}"))?;
        Ok(Self { client })
    }

    pub fn fetch_structure(&self, id: &str) -> Result<String, String> {
        let id = validate_atlas_id(id)?;
        let url = format!("{API_ROOT}/fetchPredictedStructure/{id}");
        let response = self
            .client
            .get(&url)
            .send()
            .map_err(|error| format!("ESM Atlas request failed: {error}"))?;
        response_to_pdb(response, &format!("ESM Atlas ID {id}"))
    }

    pub fn fold_sequence(&self, sequence: &str) -> Result<String, String> {
        let sequence = normalize_sequence(sequence)?;
        let url = format!("{API_ROOT}/foldSequence/v1/pdb/");
        let response = self
            .client
            .post(&url)
            .header(reqwest::header::CONTENT_TYPE, "text/plain")
            .body(sequence)
            .send()
            .map_err(|error| format!("ESMFold request failed: {error}"))?;
        response_to_pdb(response, "ESMFold prediction")
    }
}

fn response_to_pdb(
    response: reqwest::blocking::Response,
    operation: &str,
) -> Result<String, String> {
    let status = response.status();
    let text = response
        .text()
        .map_err(|error| format!("cannot read {operation} response: {error}"))?;
    if !status.is_success() {
        let detail = text.trim();
        return Err(if detail.is_empty() {
            format!("{operation} failed with HTTP {status}")
        } else {
            format!("{operation} failed with HTTP {status}: {detail}")
        });
    }
    if !text.lines().any(|line| {
        let record = line.get(..6).unwrap_or(line).trim();
        record == "ATOM" || record == "HETATM"
    }) {
        return Err(format!(
            "{operation} returned a successful response, but it was not PDB data"
        ));
    }
    Ok(text)
}

pub fn validate_atlas_id(input: &str) -> Result<String, String> {
    let id = input.trim().to_ascii_uppercase();
    let suffix = id.strip_prefix("MGYP").ok_or_else(|| {
        format!("invalid ESM Atlas ID {input:?}: identifiers must start with MGYP")
    })?;
    if suffix.is_empty() || !suffix.chars().all(|character| character.is_ascii_digit()) {
        return Err(format!(
            "invalid ESM Atlas ID {input:?}: MGYP must be followed by digits"
        ));
    }
    Ok(id)
}

pub fn normalize_sequence(input: &str) -> Result<String, String> {
    let mut sequence = String::new();
    for line in input.lines() {
        let trimmed = line.trim();
        if trimmed.starts_with('>') {
            continue;
        }
        for character in trimmed.chars() {
            if character.is_ascii_whitespace() || character.is_ascii_digit() {
                continue;
            }
            let amino_acid = character.to_ascii_uppercase();
            if !matches!(
                amino_acid,
                'A' | 'R'
                    | 'N'
                    | 'D'
                    | 'C'
                    | 'E'
                    | 'Q'
                    | 'G'
                    | 'H'
                    | 'I'
                    | 'L'
                    | 'K'
                    | 'M'
                    | 'F'
                    | 'P'
                    | 'S'
                    | 'T'
                    | 'W'
                    | 'Y'
                    | 'V'
                    | 'B'
                    | 'J'
                    | 'O'
                    | 'U'
                    | 'X'
                    | 'Z'
            ) {
                return Err(format!(
                    "invalid amino-acid character {character:?} in sequence"
                ));
            }
            sequence.push(amino_acid);
        }
    }
    if sequence.len() < 4 {
        return Err("protein sequence must contain at least 4 residues".to_owned());
    }
    if sequence.len() > 1_024 {
        return Err(format!(
            "sequence has {} residues; the public ESMFold endpoint is limited to small interactive jobs (maximum 1024 in MolX)",
            sequence.len()
        ));
    }
    Ok(sequence)
}

pub fn read_sequence_file(path: impl AsRef<Path>) -> Result<String, String> {
    let path = path.as_ref();
    let text = fs::read_to_string(path)
        .map_err(|error| format!("cannot read sequence file {}: {error}", path.display()))?;
    normalize_sequence(&text)
}

pub fn write_pdb(path: impl AsRef<Path>, pdb: &str, force: bool) -> Result<(), String> {
    let path = path.as_ref();
    if path.exists() && !force {
        return Err(format!(
            "{} already exists; pass --force to replace it",
            path.display()
        ));
    }
    fs::write(path, pdb).map_err(|error| format!("cannot write {}: {error}", path.display()))
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn validates_and_normalizes_atlas_id() {
        assert_eq!(
            validate_atlas_id("mgyp002537940442").unwrap(),
            "MGYP002537940442"
        );
        assert!(validate_atlas_id("P12345").is_err());
        assert!(validate_atlas_id("MGYP12X").is_err());
    }

    #[test]
    fn cleans_fasta_sequence() {
        assert_eq!(
            normalize_sequence(">protein one\nMKT 12\nAYX\n").unwrap(),
            "MKTAYX"
        );
        assert!(normalize_sequence("MKT*").is_err());
        assert!(normalize_sequence("MK").is_err());
    }
}