anyreader/
format.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
use crate::AnyReader;
use peekable::Peekable;
use std::fmt::{Debug, Formatter};
use std::io;
use std::io::Read;
use tracing::trace;

/// A reader that contains a detected file format.
///
/// ## Read from compressed formats
/// ```
/// # use anyreader::{AnyFormat, test::{gzip_data, tar_archive}};
/// let compressed_data = gzip_data(b"hello compressed world");
/// let mut reader = AnyFormat::from_reader(compressed_data.as_slice()).unwrap();
/// assert!(reader.kind.is_gzip());
/// assert_eq!(std::io::read_to_string(reader).unwrap(), "hello compressed world");
/// ```
///
/// ## Detect and read from compressed archive formats
/// ```
/// # use anyreader::{AnyFormat, test::{gzip_data, tar_archive}};
/// let tar_gz = gzip_data(tar_archive([("test", b"hello tar world")]));
/// let mut reader = AnyFormat::from_reader(tar_gz.as_slice()).unwrap();
/// assert!(reader.kind.is_tar());
/// let mut archive = tar::Archive::new(reader);
/// let mut entry = archive.entries().unwrap().next().unwrap().unwrap();
/// assert_eq!(std::io::read_to_string(entry).unwrap(), "hello tar world");
/// ```
pub struct AnyFormat<T: Read> {
    pub kind: FormatKind,
    reader: Peekable<AnyReader<T>>,
}

impl<T: Read> AnyFormat<T> {
    pub fn from_reader(reader: T) -> io::Result<AnyFormat<T>> {
        const MAX_PEEK_BUFFER_SIZE: usize = 262;

        let compression_reader = AnyReader::from_reader(reader)?;
        let format: FormatKind = (&compression_reader).into();
        let mut reader = Peekable::with_capacity(compression_reader, MAX_PEEK_BUFFER_SIZE);
        reader.fill_peek_buf()?;
        let buf = crate::peek_upto::<MAX_PEEK_BUFFER_SIZE>(&mut reader)?;
        trace!("peeked {} bytes", buf.len());

        let format: FormatKind = if infer::archive::is_tar(buf) {
            FormatKind::Tar
        } else if infer::archive::is_zip(buf) {
            FormatKind::Zip
        } else {
            format
        };

        trace!("format detected: {format:?}");

        Ok(AnyFormat {
            kind: format,
            reader,
        })
    }

    pub fn get_ref(&self) -> &T {
        self.reader.get_ref().1.get_ref()
    }
}

impl<T: Read> Debug for AnyFormat<T> {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("AnyFormat")
            .field("kind", &self.kind)
            .finish()
    }
}

impl<T: Read> Read for AnyFormat<T> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        self.reader.read(buf)
    }
}

/// Supported file/compression formats.
#[derive(
    Debug,
    Copy,
    Clone,
    Eq,
    PartialEq,
    Hash,
    Default,
    strum::EnumString,
    strum::Display,
    strum::EnumIs,
)]
#[strum(serialize_all = "lowercase", ascii_case_insensitive)]
pub enum FormatKind {
    /// Gzip compression
    Gzip,
    /// ZStandard compression
    Zstd,
    /// Bzip2 compression
    Bzip2,
    /// XZ compression
    Xz,
    /// Zip archive
    Zip,
    /// Tar archive. Note: this may be compressed with any of the
    /// previous compression formats (i.e. tar.gz, tar.zst, ...)
    Tar,
    /// Unknown format. This is the fallback when the format is not recognized, and
    /// the associated [AnyFormat] will read the data as-is.
    #[default]
    Unknown,
}

impl<T: Read> From<&AnyReader<T>> for FormatKind {
    /// Convert a `CompressionReader` into a `FormatKind`.
    fn from(reader: &AnyReader<T>) -> Self {
        match reader {
            AnyReader::Gzip(_) => FormatKind::Gzip,
            AnyReader::Zst(_) => FormatKind::Zstd,
            AnyReader::Bzip2(_) => FormatKind::Bzip2,
            AnyReader::Xz(_) => FormatKind::Xz,
            AnyReader::Unknown(_) => FormatKind::Unknown,
        }
    }
}

impl<T: Read> From<AnyReader<T>> for FormatKind {
    fn from(reader: AnyReader<T>) -> Self {
        (&reader).into()
    }
}