Skip to main content

fits_io/
fits_slice.rs

1use crate::fits::Fits;
2use crate::hdu::{ExtensionHDU, HDU};
3use crate::header::header::BLOCK_NUM_BYTES;
4use crate::header::{ExtensionType, Header};
5use crate::slice_ascii_table_hdu::SliceAsciiTableHDU;
6use crate::slice_bin_table_hdu::SliceBinTableHDU;
7use crate::slice_image_hdu::SliceImageHDU;
8use log::debug;
9use std::error::Error;
10use std::io::{Cursor, Seek, SeekFrom};
11use std::sync::Arc;
12
13/// A FITS file read from a buffer rather than from the filesystem.
14///
15/// This is the reader to use where there is no filesystem to read from — a file
16/// already in memory, one arriving over a network, or a build with the `fs`
17/// feature turned off.
18#[derive(Debug, Clone)]
19pub struct FitsSlice {
20    primary_hdu: SliceImageHDU,
21    extension_hdus: Vec<ExtensionHDU<Self>>,
22}
23
24impl FitsSlice {
25    /// Reads a FITS file out of `data`.
26    ///
27    /// The buffer is copied once, because the HDUs read from it for as long as
28    /// they live and cannot borrow from the caller's slice. Use
29    /// [`FitsSlice::from_vec`] to hand over a buffer instead of copying one.
30    pub fn from_slice(data: &[u8]) -> Result<Self, Box<dyn Error + Send + Sync>> {
31        Self::from_vec(data.to_vec())
32    }
33
34    /// Reads a FITS file out of `data`, taking ownership of the buffer.
35    ///
36    /// A buffer that holds gzip-compressed data is decompressed first. A FITS
37    /// file starts with `SIMPLE`, so it can never be mistaken for one that
38    /// starts with the gzip marker.
39    pub fn from_vec(data: Vec<u8>) -> Result<Self, Box<dyn Error + Send + Sync>> {
40        #[cfg(feature = "gzip")]
41        let data = if is_gzipped(&data) {
42            decompress(&data)?
43        } else {
44            data
45        };
46
47        let data = Arc::new(data);
48        let mut reader: Box<dyn crate::util::ReadSeek> =
49            Box::new(Cursor::new(crate::util::SharedBytes(Arc::clone(&data))));
50
51        let header =
52            Header::from_reader(&mut reader)?.ok_or("Could not read primary FITS header")?;
53        header.validate_primary()?;
54        debug!("Read primary header: {:?}", header);
55
56        let data_offset = header.bytes_len();
57        let primary_hdu = SliceImageHDU::new(header, Arc::clone(&data), data_offset);
58
59        let mut extension_hdus = vec![];
60        let mut offset = primary_hdu.byte_size();
61
62        loop {
63            if offset as usize >= data.len() {
64                break;
65            }
66
67            reader.seek(SeekFrom::Start(offset))?;
68
69            let Some(header) = Header::from_reader(&mut reader)? else {
70                break;
71            };
72            header.validate_extension()?;
73
74            let extension_type = header
75                .extension()
76                .ok_or("This is not a valid fits extension. Card XTENSION is missing or invalid")?;
77
78            let data_offset = offset as usize + header.bytes_len();
79
80            match extension_type {
81                ExtensionType::Image => {
82                    let hdu = SliceImageHDU::new(header, Arc::clone(&data), data_offset);
83                    offset += hdu.byte_size();
84                    extension_hdus.push(ExtensionHDU::Image(hdu));
85                }
86                // A compressed image is stored as a table, but it is an image,
87                // and presenting it as a table of opaque bytes would leave every
88                // caller to notice and unpack it themselves.
89                ExtensionType::BinTable if header.is_compressed_image() => {
90                    let hdu = SliceImageHDU::new(header, Arc::clone(&data), data_offset);
91                    offset += hdu.byte_size();
92                    extension_hdus.push(ExtensionHDU::Image(hdu));
93                }
94                ExtensionType::BinTable => {
95                    let hdu = SliceBinTableHDU::new(header, Arc::clone(&data), data_offset);
96                    offset += hdu.byte_size();
97                    extension_hdus.push(ExtensionHDU::BinTable(hdu));
98                }
99                ExtensionType::AsciiTable => {
100                    let hdu = SliceAsciiTableHDU::new(header, Arc::clone(&data), data_offset);
101                    offset += hdu.byte_size();
102                    extension_hdus.push(ExtensionHDU::AsciiTable(hdu));
103                }
104            }
105        }
106
107        Ok(Self {
108            primary_hdu,
109            extension_hdus,
110        })
111    }
112
113    /// Reads a FITS file out of `data` without blocking the async runtime.
114    ///
115    /// Parsing a large buffer is real work even though it touches no
116    /// filesystem, so it runs on a blocking worker, as
117    /// [`FsFits::open_async`](crate::fs::FsFits::open_async) does.
118    #[cfg(feature = "tokio")]
119    pub async fn from_vec_async(data: Vec<u8>) -> Result<Self, Box<dyn Error + Send + Sync>> {
120        tokio::task::spawn_blocking(move || Self::from_vec(data)).await?
121    }
122
123    /// An empty FITS file, for building one from nothing.
124    pub fn new() -> Self {
125        Self {
126            primary_hdu: SliceImageHDU::empty(),
127            extension_hdus: vec![],
128        }
129    }
130}
131
132impl Default for FitsSlice {
133    fn default() -> Self {
134        Self::new()
135    }
136}
137
138impl Fits for FitsSlice {
139    type ImageHDU = SliceImageHDU;
140    type BinTableHDU = SliceBinTableHDU;
141    type AsciiTableHDU = SliceAsciiTableHDU;
142
143    fn primary_hdu(&self) -> &Self::ImageHDU {
144        &self.primary_hdu
145    }
146
147    fn primary_hdu_mut(&mut self) -> &mut Self::ImageHDU {
148        &mut self.primary_hdu
149    }
150
151    fn extension_count(&self) -> usize {
152        self.extension_hdus.len()
153    }
154
155    fn extension_hdu(&self, index: usize) -> Option<&ExtensionHDU<Self>> {
156        self.extension_hdus.get(index)
157    }
158
159    fn extension_hdu_mut(&mut self, index: usize) -> Option<&mut ExtensionHDU<Self>> {
160        self.extension_hdus.get_mut(index)
161    }
162
163    fn extension_hdus(&self) -> impl Iterator<Item = &ExtensionHDU<Self>> {
164        self.extension_hdus.iter()
165    }
166
167    fn extension_hdus_mut(&mut self) -> impl Iterator<Item = &mut ExtensionHDU<Self>> {
168        self.extension_hdus.iter_mut()
169    }
170
171    fn push_extension(&mut self, extension: ExtensionHDU<Self>) {
172        self.extension_hdus.push(extension);
173    }
174
175    fn remove_extension(&mut self, index: usize) -> Option<ExtensionHDU<Self>> {
176        (index < self.extension_hdus.len()).then(|| self.extension_hdus.remove(index))
177    }
178
179    fn to_vec(&self) -> Result<Vec<u8>, Box<dyn Error + Send + Sync>> {
180        let mut bytes = Vec::new();
181
182        // A compressed image lives in a binary table, and the primary HDU of a
183        // FITS file cannot be one. Rather than write a file no reader would
184        // accept, the image goes into the first extension and the primary HDU
185        // is left empty — which is what `fpack` does with the same problem, and
186        // is what this crate reads back as the image it was.
187        if self.primary_hdu.header().is_compressed_image() {
188            append_hdu(&mut bytes, &empty_primary_header(), &[], 0)?;
189            append_hdu(
190                &mut bytes,
191                &self
192                    .primary_hdu
193                    .header()
194                    .conformed(Some(ExtensionType::BinTable)),
195                self.primary_hdu.data_bytes(),
196                0,
197            )?;
198        } else {
199            append_hdu(
200                &mut bytes,
201                &self.primary_hdu.header().conformed(None),
202                self.primary_hdu.data_bytes(),
203                0,
204            )?;
205        }
206
207        for extension in &self.extension_hdus {
208            match extension {
209                // A compressed image is written as the table it is stored
210                // as; XTENSION says how to read the bytes, not what they mean.
211                ExtensionHDU::Image(hdu) => append_hdu(
212                    &mut bytes,
213                    &hdu.header()
214                        .conformed(Some(extension_type_of(hdu.header()))),
215                    hdu.data_bytes(),
216                    0,
217                )?,
218                ExtensionHDU::BinTable(hdu) => append_hdu(
219                    &mut bytes,
220                    &hdu.header().conformed(Some(ExtensionType::BinTable)),
221                    hdu.data_bytes(),
222                    0,
223                )?,
224                // An ASCII table holds characters, and the standard pads it with
225                // the blanks that a character field means, not with zero bytes.
226                ExtensionHDU::AsciiTable(hdu) => append_hdu(
227                    &mut bytes,
228                    &hdu.header().conformed(Some(ExtensionType::AsciiTable)),
229                    hdu.data_bytes(),
230                    b' ',
231                )?,
232            }
233        }
234
235        Ok(bytes)
236    }
237}
238
239/// The two bytes that open every gzip stream.
240#[cfg(feature = "gzip")]
241fn is_gzipped(data: &[u8]) -> bool {
242    data.starts_with(&[0x1f, 0x8b])
243}
244
245#[cfg(feature = "gzip")]
246fn decompress(data: &[u8]) -> Result<Vec<u8>, Box<dyn Error + Send + Sync>> {
247    use std::io::Read;
248
249    let mut decoder = flate2::read::GzDecoder::new(data);
250    let mut decompressed = Vec::new();
251    decoder.read_to_end(&mut decompressed)?;
252
253    Ok(decompressed)
254}
255
256/// Appends one HDU: its header, its data, and the padding that squares the data
257/// off to a whole number of blocks.
258///
259/// The header is rendered last, because its CHECKSUM card covers the padded data
260/// as well as the header itself.
261fn append_hdu(
262    bytes: &mut Vec<u8>,
263    header: &Header,
264    data: &[u8],
265    padding: u8,
266) -> Result<(), Box<dyn Error + Send + Sync>> {
267    header.validate_against_data(data.len())?;
268
269    let mut data = data.to_vec();
270    let overhang = data.len() % BLOCK_NUM_BYTES;
271    if overhang != 0 {
272        data.resize(data.len() + BLOCK_NUM_BYTES - overhang, padding);
273    }
274
275    bytes.extend_from_slice(&header.checksummed_bytes(&data));
276    bytes.extend_from_slice(&data);
277
278    Ok(())
279}
280
281/// Which kind of extension an image HDU is written as.
282///
283/// A tile-compressed image lives in a binary table, and a reader finds it by its
284/// XTENSION before it ever looks at the `Z` keywords that say it is an image.
285fn extension_type_of(header: &Header) -> ExtensionType {
286    if header.is_compressed_image() {
287        ExtensionType::BinTable
288    } else {
289        ExtensionType::Image
290    }
291}
292
293/// The header of a primary HDU holding no data, for a file whose image had to
294/// move into an extension.
295fn empty_primary_header() -> Header {
296    let mut header = Header::default();
297
298    // Without EXTEND, a reader is entitled to stop at the primary HDU and never
299    // look for the extension the image is in.
300    let _ = header.set_card(crate::header::card_keys::EXTEND, true);
301
302    header.conformed(None)
303}