winsfs_core/sfs/io/
plain_text.rs

1//! Reading and writing for the plain text SFS format.
2//!
3//! The plain text format is a simple format consisting of two lines.
4//! The first line contains a header line `#SHAPE=<[shape]>`, where `[shape]`
5//! is a `/`-separated representation of the shape of the SFS. The next line
6//! gives the SFS in flat, row-major order separated by a single space.
7//!
8//! In other words, the plain text format is like the format output by realSFS,
9//! except with the addition of a header line so that the SFS can be read without
10//! passing the shape separately.
11
12use std::{error::Error, fmt, fs::File, io, path::Path, str::FromStr};
13
14use crate::sfs::{
15    generics::{DynShape, Normalisation, Shape},
16    DynUSfs, SfsBase,
17};
18
19/// Parses an SFS in plain text format from the raw, flat text representation.
20///
21/// `s` is assumed to not contain the header.
22fn parse_sfs(s: &str, shape: DynShape) -> io::Result<DynUSfs> {
23    s.split_ascii_whitespace()
24        .map(f64::from_str)
25        .collect::<Result<Vec<_>, _>>()
26        .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))
27        .and_then(|vec| {
28            DynUSfs::from_vec_shape(vec, shape)
29                .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))
30        })
31}
32
33/// Reads an SFS in plain text format from a reader.
34///
35/// The stream is assumed to be positioned at the start.
36pub fn read_sfs<R>(reader: &mut R) -> io::Result<DynUSfs>
37where
38    R: io::BufRead,
39{
40    let header = Header::read(reader)?;
41
42    let mut buf = String::new();
43    let _bytes_read = reader.read_to_string(&mut buf)?;
44
45    parse_sfs(&buf, header.shape)
46}
47
48/// Reads an SFS in plain text format from a file path.
49pub fn read_sfs_from_path<P>(path: P) -> io::Result<DynUSfs>
50where
51    P: AsRef<Path>,
52{
53    let mut reader = File::open(path).map(io::BufReader::new)?;
54    read_sfs(&mut reader)
55}
56
57/// Writes an SFS in plain text format to a writer.
58pub fn write_sfs<W, S, N>(writer: &mut W, sfs: &SfsBase<S, N>) -> io::Result<()>
59where
60    W: io::Write,
61    S: Shape,
62    N: Normalisation,
63{
64    let header = Header::new(sfs.shape().as_ref().to_vec().into_boxed_slice());
65    header.write(writer)?;
66
67    writeln!(writer, "{}", sfs.format_flat(" ", 6))
68}
69
70/// Writes an SFS in plain text format to a file path.
71///
72/// If the file already exists, it will be overwritten.
73pub fn write_sfs_to_path<P, S, N>(path: P, sfs: &SfsBase<S, N>) -> io::Result<()>
74where
75    P: AsRef<Path>,
76    S: Shape,
77    N: Normalisation,
78{
79    let mut writer = File::create(path)?;
80    write_sfs(&mut writer, sfs)
81}
82
83/// A plain text SFS header.
84#[derive(Clone, Debug)]
85struct Header {
86    shape: DynShape,
87}
88
89impl Header {
90    /// Creates a new header.
91    pub fn new(shape: DynShape) -> Self {
92        Self { shape }
93    }
94
95    /// Reads a header from a reader.
96    ///
97    /// Assumes the stream is positioned immediately in front of the header.
98    pub fn read<R>(reader: &mut R) -> io::Result<Self>
99    where
100        R: io::BufRead,
101    {
102        let mut buf = String::new();
103
104        reader.read_line(&mut buf)?;
105
106        Self::from_str(&buf).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))
107    }
108
109    /// Writes a header to a stream.
110    pub fn write<W>(&self, writer: &mut W) -> io::Result<()>
111    where
112        W: io::Write,
113    {
114        writeln!(writer, "{self}")
115    }
116}
117
118impl fmt::Display for Header {
119    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
120        let shape_fmt = self
121            .shape
122            .iter()
123            .map(|x| x.to_string())
124            .collect::<Vec<_>>()
125            .join("/");
126
127        write!(f, "#SHAPE=<{shape_fmt}>")
128    }
129}
130
131impl FromStr for Header {
132    type Err = ParseHeaderError;
133
134    fn from_str(s: &str) -> Result<Self, Self::Err> {
135        s.trim_start_matches(|c: char| !c.is_numeric())
136            .trim_end_matches(|c: char| !c.is_numeric())
137            .split('/')
138            .map(usize::from_str)
139            .collect::<Result<Vec<_>, _>>()
140            .map_err(|_| ParseHeaderError(String::from(s)))
141            .map(Vec::into_boxed_slice)
142            .map(Header::new)
143    }
144}
145
146/// An error associated with parsing the plain text format header.
147#[derive(Debug)]
148pub struct ParseHeaderError(String);
149
150impl fmt::Display for ParseHeaderError {
151    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
152        write!(f, "failed to parse '{}' as plain SFS format header", self.0)
153    }
154}
155
156impl Error for ParseHeaderError {}
157
158#[cfg(test)]
159mod tests {
160    use super::*;
161
162    use crate::{sfs1d, sfs2d};
163
164    #[test]
165    fn test_parse_header() {
166        assert_eq!(Header::from_str("#SHAPE=<3>").unwrap().shape.as_ref(), [3]);
167        assert_eq!(
168            Header::from_str("#SHAPE=<11/13>").unwrap().shape.as_ref(),
169            &[11, 13]
170        );
171    }
172
173    #[test]
174    fn test_display_header() {
175        assert_eq!(Header::new(Box::new([25])).to_string(), "#SHAPE=<25>");
176        assert_eq!(Header::new(Box::new([7, 9])).to_string(), "#SHAPE=<7/9>");
177    }
178
179    #[test]
180    fn test_read_1d() -> io::Result<()> {
181        let src = b"#SHAPE=<3>\n0.0 1.0 2.0\n";
182
183        assert_eq!(read_sfs(&mut &src[..])?, DynUSfs::from(sfs1d![0., 1., 2.]));
184
185        Ok(())
186    }
187
188    #[test]
189    fn test_read_2d() -> io::Result<()> {
190        let src = b"#SHAPE=<2/3>\n0.0 1.0 2.0 3.0 4.0 5.0\n";
191
192        assert_eq!(
193            read_sfs(&mut &src[..])?,
194            DynUSfs::from(sfs2d![[0., 1., 2.], [3., 4., 5.]])
195        );
196
197        Ok(())
198    }
199
200    #[test]
201    fn test_write_1d() -> io::Result<()> {
202        let mut dest = Vec::new();
203        write_sfs(&mut dest, &sfs1d![0., 1., 2.])?;
204
205        assert_eq!(dest, b"#SHAPE=<3>\n0.000000 1.000000 2.000000\n");
206
207        Ok(())
208    }
209
210    #[test]
211    fn test_write_2d() -> io::Result<()> {
212        let mut dest = Vec::new();
213        write_sfs(&mut dest, &sfs2d![[0., 1., 2.], [3., 4., 5.]])?;
214
215        assert_eq!(
216            dest,
217            b"#SHAPE=<2/3>\n0.000000 1.000000 2.000000 3.000000 4.000000 5.000000\n",
218        );
219
220        Ok(())
221    }
222}