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}
31
32impl Format {
33    /// The formats a caller may name, excluding automatic detection.
34    pub const NAMED: [Self; 7] = [
35        Self::Mmcif,
36        Self::Pdbml,
37        Self::BinaryCif,
38        Self::Mmtf,
39        Self::Pdbqt,
40        Self::Pqr,
41        Self::Pdb,
42    ];
43
44    /// The format's short name.
45    #[must_use]
46    pub const fn name(self) -> &'static str {
47        match self {
48            Self::Auto => "auto",
49            Self::Mmcif => "mmcif",
50            Self::Pdbml => "pdbml",
51            Self::BinaryCif => "bcif",
52            Self::Mmtf => "mmtf",
53            Self::Pdb => "pdb",
54            Self::Pqr => "pqr",
55            Self::Pdbqt => "pdbqt",
56        }
57    }
58
59    /// The file suffixes conventionally used for this format.
60    #[must_use]
61    pub const fn extensions(self) -> &'static [&'static str] {
62        match self {
63            Self::Auto => &[],
64            Self::Mmcif => &["cif", "mmcif"],
65            Self::Pdbml => &["xml", "pdbml"],
66            Self::BinaryCif => &["bcif"],
67            Self::Mmtf => &["mmtf"],
68            Self::Pdb => &["pdb", "ent"],
69            Self::Pqr => &["pqr"],
70            Self::Pdbqt => &["pdbqt"],
71        }
72    }
73
74    /// Parses a format name without allocating a normalized copy.
75    #[must_use]
76    pub fn parse(name: &str) -> Option<Self> {
77        if name.eq_ignore_ascii_case("auto") {
78            Some(Self::Auto)
79        } else if name.eq_ignore_ascii_case("mmcif") || name.eq_ignore_ascii_case("cif") {
80            Some(Self::Mmcif)
81        } else if name.eq_ignore_ascii_case("pdbml") || name.eq_ignore_ascii_case("xml") {
82            Some(Self::Pdbml)
83        } else if name.eq_ignore_ascii_case("bcif") || name.eq_ignore_ascii_case("binarycif") {
84            Some(Self::BinaryCif)
85        } else if name.eq_ignore_ascii_case("mmtf") {
86            Some(Self::Mmtf)
87        } else if name.eq_ignore_ascii_case("pdb") || name.eq_ignore_ascii_case("ent") {
88            Some(Self::Pdb)
89        } else if name.eq_ignore_ascii_case("pqr") {
90            Some(Self::Pqr)
91        } else if name.eq_ignore_ascii_case("pdbqt") {
92            Some(Self::Pdbqt)
93        } else {
94            None
95        }
96    }
97
98    /// Returns true when this format's content is recognisable in `bytes`.
99    #[must_use]
100    pub fn recognises(self, bytes: &[u8]) -> bool {
101        match self {
102            Self::Auto | Self::Pqr => false,
103            Self::Mmcif => first_lines(bytes, 8).any(|line| line.starts_with(b"data_")),
104            Self::Pdbml => contains_in_probe(bytes, b"PDBx:datablock", PDBML_PROBE_WINDOWS),
105            Self::BinaryCif => recognises_messagepack_map(bytes, b"dataBlocks"),
106            Self::Mmtf => recognises_messagepack_map(bytes, b"mmtfVersion"),
107            Self::Pdbqt => recognises_pdbqt(bytes),
108            Self::Pdb => recognises_pdb(bytes),
109        }
110    }
111
112    /// Decides the format of `input`, given what the caller asked for and the
113    /// name the bytes came from.
114    ///
115    /// # Errors
116    ///
117    /// Returns a diagnostic when neither content nor the optional name identifies
118    /// a supported format.
119    pub fn detect(
120        requested: Self,
121        input: &InputBuffer,
122        name: Option<&str>,
123    ) -> Result<Self, Diagnostic> {
124        if requested != Self::Auto {
125            return Ok(requested);
126        }
127
128        let bytes = input.as_bytes();
129        for format in [
130            Self::Mmcif,
131            Self::Pdbml,
132            Self::BinaryCif,
133            Self::Mmtf,
134            Self::Pdbqt,
135        ] {
136            if format.recognises(bytes) {
137                return Ok(format);
138            }
139        }
140
141        let named = name.and_then(Self::from_name);
142        if recognises_pdb(bytes) {
143            if let Some(format @ (Self::Pqr | Self::Pdbqt)) = named {
144                return Ok(format);
145            }
146            return Ok(Self::Pdb);
147        }
148        if let Some(format) = named {
149            return Ok(format);
150        }
151
152        let Some(name) = name else {
153            return Err(Diagnostic::new(Code::E1001)
154                .with_context("tried", "content, then file name")
155                .with_context("name", "(none)"));
156        };
157        Err(Diagnostic::new(Code::E1001)
158            .with_context("tried", "content, then file name")
159            .with_context("name", name))
160    }
161
162    /// The format a file name suggests, ignoring any compression suffix.
163    #[must_use]
164    pub fn from_name(name: &str) -> Option<Self> {
165        let mut suffixes = name.rsplit('.');
166        let mut candidate = suffixes.next()?;
167        if is_compression_suffix(candidate) {
168            candidate = suffixes.next()?;
169        }
170        Self::NAMED.into_iter().find(|format| {
171            format
172                .extensions()
173                .iter()
174                .any(|extension| candidate.eq_ignore_ascii_case(extension))
175        })
176    }
177}
178
179fn is_compression_suffix(suffix: &str) -> bool {
180    suffix.eq_ignore_ascii_case("gz")
181        || suffix.eq_ignore_ascii_case("zst")
182        || suffix.eq_ignore_ascii_case("bz2")
183        || suffix.eq_ignore_ascii_case("xz")
184}
185
186fn is_messagepack_map_prefix(byte: u8) -> bool {
187    (0x80..=0x8f).contains(&byte) || matches!(byte, 0xde | 0xdf)
188}
189
190fn recognises_messagepack_map(bytes: &[u8], key: &[u8]) -> bool {
191    bytes
192        .first()
193        .copied()
194        .is_some_and(is_messagepack_map_prefix)
195        && contains_in_probe(bytes, key, MESSAGEPACK_PROBE_WINDOWS)
196}
197
198fn contains_in_probe(bytes: &[u8], needle: &[u8], windows: usize) -> bool {
199    if needle.is_empty() {
200        return true;
201    }
202    bytes
203        .windows(needle.len())
204        .take(windows)
205        .any(|window| window == needle)
206}
207
208fn recognises_pdb(bytes: &[u8]) -> bool {
209    first_lines(bytes, 64).any(|line| {
210        line.starts_with(b"ATOM  ")
211            || line.starts_with(b"HETATM")
212            || line.starts_with(b"HEADER")
213            || line.starts_with(b"CRYST1")
214            || line.starts_with(b"MODEL ")
215    })
216}
217
218fn recognises_pdbqt(bytes: &[u8]) -> bool {
219    first_lines(bytes, 64).any(recognises_pdbqt_line)
220}
221
222fn recognises_pdbqt_line(line: &[u8]) -> bool {
223    if matches!(line, b"ROOT" | b"ENDROOT") || line.starts_with(b"TORSDOF") {
224        return true;
225    }
226    if !(line.starts_with(b"ATOM  ") || line.starts_with(b"HETATM")) {
227        return false;
228    }
229    let charge = line
230        .get(70..76)
231        .and_then(|field| str::from_utf8(field).ok());
232    let atom_type = line.get(77..).and_then(|field| str::from_utf8(field).ok());
233    charge.is_some_and(|field| field.trim().parse::<f64>().is_ok())
234        && atom_type.is_some_and(|field| !field.trim().is_empty())
235}
236
237fn first_lines(bytes: &[u8], count: usize) -> impl Iterator<Item = &[u8]> {
238    bytes
239        .split(|byte| *byte == b'\n')
240        .take(count)
241        .map(strip_carriage_return)
242}
243
244fn strip_carriage_return(line: &[u8]) -> &[u8] {
245    if line.last().copied() == Some(b'\r') {
246        &line[..line.len() - 1]
247    } else {
248        line
249    }
250}