ferritin_core/io/
io.rs

1use crate::AtomCollection;
2use crate::io::cif;
3use crate::io::pdb;
4use anyhow::{Context, Result};
5use std::path::Path;
6
7//
8pub fn load_structure<P: AsRef<Path>>(file_path: P) -> Result<AtomCollection> {
9    let path = file_path.as_ref();
10    let extension = path
11        .extension()
12        .and_then(|ext| ext.to_str())
13        .ok_or_else(|| anyhow::anyhow!("File has no extension"))?
14        .to_lowercase();
15
16    let mut ac = match extension.as_str() {
17        "pdb" => pdb::PDBFile::read(path)
18            .context("Failed to read PDB file")?
19            .parse_to_atom_collection()
20            .context("Failed to parse PDB file to atom collection")?,
21        "cif" => cif::CIFFile::read(path)?.parse_to_atom_collection()?,
22        _ => return Err(anyhow::anyhow!("Unsupported file extension: {}", extension)),
23    };
24
25    ac.connect_via_residue_names();
26    Ok(ac)
27}
28
29pub fn load_structure_from_string(content: &str, filetype: &str) -> Result<AtomCollection> {
30    let filetype = filetype.to_lowercase();
31
32    let mut ac = match filetype.as_str() {
33        "pdb" => pdb::PDBFile::new_from_string(content.to_string())
34            .context("Failed to read PDB from string")?
35            .parse_to_atom_collection()
36            .context("Failed to parse PDB string to atom collection")?,
37        "cif" => cif::CIFFile::new(content.to_string())
38            .context("Failed to read CIF from string")?
39            .parse_to_atom_collection()
40            .context("Failed to parse CIF string to atom collection")?,
41        _ => return Err(anyhow::anyhow!("Unsupported file type: {}", filetype)),
42    };
43
44    ac.connect_via_residue_names();
45    Ok(ac)
46}
47
48#[cfg(test)]
49mod tests {
50    use super::*;
51    use ferritin_test_data::TestFile;
52
53    #[test]
54    fn test_load_structure_cif() {
55        let (file_path, _handle) = TestFile::protein_01().create_temp().unwrap();
56        let result = load_structure(&file_path);
57        assert!(
58            result.is_ok(),
59            "Failed to load CIF file: {:?}",
60            result.err()
61        );
62        let atom_collection = result.unwrap();
63        assert_eq!(atom_collection.get_size(), 1413);
64        // This includes Water Molecules
65        let max_resid = atom_collection.get_resids().iter().max().unwrap_or(&0);
66        assert_eq!(*max_resid, 338);
67    }
68
69    #[test]
70    fn test_load_structure_pdb() {
71        let (file_path, _handle) = TestFile::protein_02().create_temp().unwrap();
72        let result = load_structure(&file_path);
73        assert!(
74            result.is_ok(),
75            "Failed to load PDB file: {:?}",
76            result.err()
77        );
78        let atom_collection = result.unwrap();
79        assert_eq!(atom_collection.get_size(), 1356);
80        let max_resid = atom_collection.get_resids().iter().max().unwrap_or(&0);
81        assert_eq!(*max_resid, 176);
82    }
83}