1use crate::AtomCollection;
2use crate::io::cif;
3use crate::io::pdb;
4use anyhow::{Context, Result};
5use std::path::Path;
6
7pub 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 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}