Skip to main content

molframe_core/io/format/
detect.rs

1//! Content and suffix based format recognition.
2
3use super::super::input::InputBuffer;
4use crate::diagnostic::{Code, Diagnostic};
5
6const PDBML_PROBE_WINDOWS: usize = 2_048;
7const MESSAGEPACK_PROBE_WINDOWS: usize = 512;
8
9/// A format molframe can read or write.
10#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug, Default)]
11#[non_exhaustive]
12pub enum Format {
13    /// Decide from the content, then from the file name.
14    #[default]
15    Auto,
16    /// The archive's structured text format.
17    Mmcif,
18    /// XML serialization of the `PDBx` data model.
19    Pdbml,
20    /// The archive's `MessagePack`-based binary format.
21    BinaryCif,
22    /// The deprecated but still encountered macromolecular transmission format.
23    Mmtf,
24    /// The legacy fixed-column format.
25    Pdb,
26    /// PDB-shaped coordinates carrying partial charge and radius.
27    Pqr,
28    /// `AutoDock`'s PDB-shaped coordinates carrying charge and atom type.
29    Pdbqt,
30    /// MDL structure-data file: one or more `CTfile` records.
31    Sdf,
32    /// Tripos MOL2 molecule records.
33    Mol2,
34    /// Core-CIF small-molecule crystal data with fractional coordinates.
35    SmallCif,
36}
37
38impl Format {
39    /// The formats a caller may name, excluding automatic detection.
40    pub const NAMED: [Self; 10] = [
41        Self::Mmcif,
42        Self::Pdbml,
43        Self::BinaryCif,
44        Self::Mmtf,
45        Self::Pdbqt,
46        Self::Pqr,
47        Self::Pdb,
48        Self::Sdf,
49        Self::Mol2,
50        Self::SmallCif,
51    ];
52
53    /// The format's short name.
54    #[must_use]
55    pub const fn name(self) -> &'static str {
56        match self {
57            Self::Auto => "auto",
58            Self::Mmcif => "mmcif",
59            Self::Pdbml => "pdbml",
60            Self::BinaryCif => "bcif",
61            Self::Mmtf => "mmtf",
62            Self::Pdb => "pdb",
63            Self::Pqr => "pqr",
64            Self::Pdbqt => "pdbqt",
65            Self::Sdf => "sdf",
66            Self::Mol2 => "mol2",
67            Self::SmallCif => "smallcif",
68        }
69    }
70
71    /// The file suffixes conventionally used for this format.
72    #[must_use]
73    pub const fn extensions(self) -> &'static [&'static str] {
74        match self {
75            Self::Auto | Self::SmallCif => &[],
76            Self::Mmcif => &["cif", "mmcif"],
77            Self::Pdbml => &["xml", "pdbml"],
78            Self::BinaryCif => &["bcif"],
79            Self::Mmtf => &["mmtf"],
80            Self::Pdb => &["pdb", "ent"],
81            Self::Pqr => &["pqr"],
82            Self::Pdbqt => &["pdbqt"],
83            Self::Sdf => &["sdf", "mol"],
84            Self::Mol2 => &["mol2"],
85        }
86    }
87
88    /// Parses a format name without allocating a normalized copy.
89    #[must_use]
90    pub fn parse(name: &str) -> Option<Self> {
91        if name.eq_ignore_ascii_case("auto") {
92            Some(Self::Auto)
93        } else if name.eq_ignore_ascii_case("mmcif") || name.eq_ignore_ascii_case("cif") {
94            Some(Self::Mmcif)
95        } else if name.eq_ignore_ascii_case("pdbml") || name.eq_ignore_ascii_case("xml") {
96            Some(Self::Pdbml)
97        } else if name.eq_ignore_ascii_case("bcif") || name.eq_ignore_ascii_case("binarycif") {
98            Some(Self::BinaryCif)
99        } else if name.eq_ignore_ascii_case("mmtf") {
100            Some(Self::Mmtf)
101        } else if name.eq_ignore_ascii_case("pdb") || name.eq_ignore_ascii_case("ent") {
102            Some(Self::Pdb)
103        } else if name.eq_ignore_ascii_case("pqr") {
104            Some(Self::Pqr)
105        } else if name.eq_ignore_ascii_case("pdbqt") {
106            Some(Self::Pdbqt)
107        } else if name.eq_ignore_ascii_case("sdf") || name.eq_ignore_ascii_case("mol") {
108            Some(Self::Sdf)
109        } else if name.eq_ignore_ascii_case("mol2") {
110            Some(Self::Mol2)
111        } else if name.eq_ignore_ascii_case("smallcif") {
112            Some(Self::SmallCif)
113        } else {
114            None
115        }
116    }
117
118    /// Returns true when this format's content is recognisable in `bytes`.
119    #[must_use]
120    pub fn recognises(self, bytes: &[u8]) -> bool {
121        match self {
122            Self::Auto | Self::Pqr => false,
123            Self::Mmcif => first_lines(bytes, 8).any(|line| line.starts_with(b"data_")),
124            Self::Pdbml => contains_in_probe(bytes, b"PDBx:datablock", PDBML_PROBE_WINDOWS),
125            Self::BinaryCif => recognises_messagepack_map(bytes, b"dataBlocks"),
126            Self::Mmtf => recognises_messagepack_map(bytes, b"mmtfVersion"),
127            Self::Pdbqt => recognises_pdbqt(bytes),
128            Self::Pdb => recognises_pdb(bytes),
129            Self::Sdf => recognises_sdf(bytes),
130            Self::Mol2 => recognises_mol2(bytes),
131            Self::SmallCif => recognises_small_cif(bytes),
132        }
133    }
134
135    /// Decides the format of `input`, given what the caller asked for and the
136    /// name the bytes came from.
137    ///
138    /// # Errors
139    ///
140    /// Returns a diagnostic when neither content nor the optional name identifies
141    /// a supported format.
142    pub fn detect(
143        requested: Self,
144        input: &InputBuffer,
145        name: Option<&str>,
146    ) -> Result<Self, Diagnostic> {
147        if requested != Self::Auto {
148            return Ok(requested);
149        }
150
151        let bytes = input.as_bytes();
152        for format in [
153            Self::SmallCif,
154            Self::Mmcif,
155            Self::Pdbml,
156            Self::BinaryCif,
157            Self::Mmtf,
158            Self::Pdbqt,
159            Self::Sdf,
160            Self::Mol2,
161        ] {
162            if format.recognises(bytes) {
163                return Ok(format);
164            }
165        }
166
167        let named = name.and_then(Self::from_name);
168        if recognises_pdb(bytes) {
169            if let Some(format @ (Self::Pqr | Self::Pdbqt)) = named {
170                return Ok(format);
171            }
172            return Ok(Self::Pdb);
173        }
174        if let Some(format) = named {
175            return Ok(format);
176        }
177
178        let Some(name) = name else {
179            return Err(Diagnostic::new(Code::E1001)
180                .with_context("tried", "content, then file name")
181                .with_context("name", "(none)"));
182        };
183        Err(Diagnostic::new(Code::E1001)
184            .with_context("tried", "content, then file name")
185            .with_context("name", name))
186    }
187
188    /// The format a file name suggests, ignoring any compression suffix.
189    #[must_use]
190    pub fn from_name(name: &str) -> Option<Self> {
191        let mut suffixes = name.rsplit('.');
192        let mut candidate = suffixes.next()?;
193        if is_compression_suffix(candidate) {
194            candidate = suffixes.next()?;
195        }
196        Self::from_extension(candidate)
197    }
198
199    /// The format a bare extension names (`pdb`, `cif`, `bcif`, ...), compared
200    /// without regard to case.
201    #[must_use]
202    pub fn from_extension(extension: &str) -> Option<Self> {
203        Self::NAMED.into_iter().find(|format| {
204            format
205                .extensions()
206                .iter()
207                .any(|known| extension.eq_ignore_ascii_case(known))
208        })
209    }
210}
211
212fn is_compression_suffix(suffix: &str) -> bool {
213    suffix.eq_ignore_ascii_case("gz")
214        || suffix.eq_ignore_ascii_case("zst")
215        || suffix.eq_ignore_ascii_case("bz2")
216        || suffix.eq_ignore_ascii_case("xz")
217}
218
219fn is_messagepack_map_prefix(byte: u8) -> bool {
220    (0x80..=0x8f).contains(&byte) || matches!(byte, 0xde | 0xdf)
221}
222
223fn recognises_messagepack_map(bytes: &[u8], key: &[u8]) -> bool {
224    bytes
225        .first()
226        .copied()
227        .is_some_and(is_messagepack_map_prefix)
228        && contains_in_probe(bytes, key, MESSAGEPACK_PROBE_WINDOWS)
229}
230
231fn contains_in_probe(bytes: &[u8], needle: &[u8], windows: usize) -> bool {
232    if needle.is_empty() {
233        return true;
234    }
235    bytes
236        .windows(needle.len())
237        .take(windows)
238        .any(|window| window == needle)
239}
240
241fn recognises_pdb(bytes: &[u8]) -> bool {
242    first_lines(bytes, 64).any(|line| {
243        line.starts_with(b"ATOM  ")
244            || line.starts_with(b"HETATM")
245            || line.starts_with(b"HEADER")
246            || line.starts_with(b"CRYST1")
247            || line.starts_with(b"MODEL ")
248    })
249}
250
251fn recognises_pdbqt(bytes: &[u8]) -> bool {
252    first_lines(bytes, 64).any(recognises_pdbqt_line)
253}
254
255fn recognises_pdbqt_line(line: &[u8]) -> bool {
256    if matches!(line, b"ROOT" | b"ENDROOT") || line.starts_with(b"TORSDOF") {
257        return true;
258    }
259    if !(line.starts_with(b"ATOM  ") || line.starts_with(b"HETATM")) {
260        return false;
261    }
262    let charge = line
263        .get(70..76)
264        .and_then(|field| str::from_utf8(field).ok());
265    let atom_type = line.get(77..).and_then(|field| str::from_utf8(field).ok());
266    charge.is_some_and(|field| field.trim().parse::<f64>().is_ok())
267        && atom_type.is_some_and(|field| !field.trim().is_empty())
268}
269
270fn recognises_sdf(bytes: &[u8]) -> bool {
271    let mut lines = first_lines(bytes, 64);
272    if lines.any(|line| line == b"$$$$") {
273        return true;
274    }
275    first_lines(bytes, 4)
276        .nth(3)
277        .is_some_and(|line| line.ends_with(b"V2000") || line.ends_with(b"V3000"))
278}
279
280fn recognises_mol2(bytes: &[u8]) -> bool {
281    first_lines(bytes, 64).any(|line| line.starts_with(b"@<TRIPOS>MOLECULE"))
282}
283
284/// A core-CIF block: fractional atom sites and no Cartesian ones.
285fn recognises_small_cif(bytes: &[u8]) -> bool {
286    let mut data_block = false;
287    let mut fractional = false;
288    for line in first_lines(bytes, 256) {
289        data_block |= line.starts_with(b"data_");
290        fractional |=
291            line.starts_with(b"_atom_site_fract_x") || line.starts_with(b"_atom_site.fract_x");
292        if line.starts_with(b"_atom_site.Cartn_x") {
293            return false;
294        }
295    }
296    data_block && fractional
297}
298
299fn first_lines(bytes: &[u8], count: usize) -> impl Iterator<Item = &[u8]> {
300    bytes
301        .split(|byte| *byte == b'\n')
302        .take(count)
303        .map(strip_carriage_return)
304}
305
306fn strip_carriage_return(line: &[u8]) -> &[u8] {
307    if line.last().copied() == Some(b'\r') {
308        &line[..line.len() - 1]
309    } else {
310        line
311    }
312}