python_pkginfo/
distribution.rs

1use std::fmt;
2use std::io::{BufReader, Read};
3use std::path::Path;
4use std::str::FromStr;
5
6#[cfg(feature = "bzip2")]
7use bzip2::read::BzDecoder;
8use flate2::read::GzDecoder;
9#[cfg(feature = "xz")]
10use xz::bufread::XzDecoder;
11#[cfg(feature = "xz")]
12use xz::stream::Stream as XzStream;
13use zip::ZipArchive;
14
15use crate::{Error, Metadata};
16
17/// Python package distribution type
18#[derive(Debug, Clone, Copy, PartialEq, Eq)]
19pub enum DistributionType {
20    /// Source distribution
21    SDist,
22    /// Binary distribution egg format
23    Egg,
24    /// Binary distribution wheel format
25    Wheel,
26}
27
28#[derive(Debug, Clone, Copy)]
29enum SDistType {
30    Zip,
31    GzTar,
32    #[cfg(feature = "deprecated-formats")]
33    Tar,
34    #[cfg(feature = "bzip2")]
35    BzTar,
36    #[cfg(feature = "xz")]
37    XzTar,
38}
39
40/// Python package distribution
41#[derive(Debug, Clone)]
42pub struct Distribution {
43    dist_type: DistributionType,
44    metadata: Metadata,
45    python_version: String,
46}
47
48impl fmt::Display for DistributionType {
49    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
50        match self {
51            DistributionType::SDist => write!(f, "sdist"),
52            DistributionType::Egg => write!(f, "bdist_egg"),
53            DistributionType::Wheel => write!(f, "bdist_wheel"),
54        }
55    }
56}
57
58impl FromStr for SDistType {
59    type Err = Error;
60
61    fn from_str(s: &str) -> Result<Self, Self::Err> {
62        let dist_type = match s {
63            "zip" => SDistType::Zip,
64            "gz" | "tgz" => SDistType::GzTar,
65            #[cfg(feature = "deprecated-formats")]
66            "tar" => SDistType::Tar,
67            #[cfg(feature = "bzip2")]
68            "bz2" | "tbz" => SDistType::BzTar,
69            #[cfg(feature = "xz")]
70            "lz" | "lzma" | "tlz" | "txz" | "xz" => SDistType::XzTar,
71            _ => return Err(Error::UnknownDistributionType),
72        };
73        Ok(dist_type)
74    }
75}
76
77impl Distribution {
78    /// Open and parse a distribution from `path`
79    pub fn new(path: impl AsRef<Path>) -> Result<Self, Error> {
80        let path = path.as_ref();
81        let ext = path
82            .extension()
83            .and_then(|ext| ext.to_str())
84            .ok_or(Error::UnknownDistributionType)?;
85
86        Ok(if let Ok(sdist_type) = ext.parse() {
87            Self {
88                dist_type: DistributionType::SDist,
89                metadata: Self::parse_sdist(path, sdist_type)?,
90                python_version: "source".to_string(),
91            }
92        } else {
93            match ext {
94                "egg" => {
95                    let parts: Vec<&str> = path
96                        .file_stem()
97                        .unwrap()
98                        .to_str()
99                        .unwrap()
100                        .split('-')
101                        .collect();
102                    let python_version = match parts.as_slice() {
103                        [_name, _version, py_ver] => py_ver,
104                        _ => "any",
105                    };
106                    Self {
107                        dist_type: DistributionType::Egg,
108                        metadata: Self::parse_egg(path)?,
109                        python_version: python_version.to_string(),
110                    }
111                }
112                "whl" => {
113                    let parts: Vec<&str> = path
114                        .file_stem()
115                        .unwrap()
116                        .to_str()
117                        .unwrap()
118                        .split('-')
119                        .collect();
120                    let python_version = match parts.as_slice() {
121                        [_name, _version, py_ver, _abi_tag, _plat_tag] => py_ver,
122                        _ => "any",
123                    };
124                    Self {
125                        dist_type: DistributionType::Wheel,
126                        metadata: Self::parse_wheel(path)?,
127                        python_version: python_version.to_string(),
128                    }
129                }
130                _ => return Err(Error::UnknownDistributionType),
131            }
132        })
133    }
134
135    /// Returns distribution type
136    pub fn r#type(&self) -> DistributionType {
137        self.dist_type
138    }
139
140    /// Returns distribution metadata
141    pub fn metadata(&self) -> &Metadata {
142        &self.metadata
143    }
144
145    /// Returns the supported Python version tag
146    ///
147    /// For source distributions the version tag is always `source`
148    pub fn python_version(&self) -> &str {
149        &self.python_version
150    }
151
152    fn parse_sdist(path: &Path, sdist_type: SDistType) -> Result<Metadata, Error> {
153        match sdist_type {
154            SDistType::Zip => Self::parse_zip(path, "PKG-INFO"),
155            SDistType::GzTar => {
156                Self::parse_tar(GzDecoder::new(BufReader::new(fs_err::File::open(path)?)))
157            }
158            #[cfg(feature = "deprecated-formats")]
159            SDistType::Tar => Self::parse_tar(BufReader::new(fs_err::File::open(path)?)),
160            #[cfg(feature = "bzip2")]
161            SDistType::BzTar => {
162                Self::parse_tar(BzDecoder::new(BufReader::new(fs_err::File::open(path)?)))
163            }
164            #[cfg(feature = "xz")]
165            SDistType::XzTar => Self::parse_tar(XzDecoder::new_stream(
166                BufReader::new(fs_err::File::open(path)?),
167                XzStream::new_auto_decoder(u64::MAX, 0).unwrap(),
168            )),
169        }
170    }
171
172    fn parse_egg(path: &Path) -> Result<Metadata, Error> {
173        Self::parse_zip(path, "EGG-INFO/PKG-INFO")
174    }
175
176    fn parse_wheel(path: &Path) -> Result<Metadata, Error> {
177        Self::parse_zip(path, ".dist-info/METADATA")
178    }
179
180    fn parse_tar<R: Read>(reader: R) -> Result<Metadata, Error> {
181        let mut reader = tar::Archive::new(reader);
182        let metadata_file = reader
183            .entries()?
184            .map(|entry| -> Result<_, Error> {
185                let entry = entry?;
186                if entry.path()?.ends_with("PKG-INFO") {
187                    Ok(Some(entry))
188                } else {
189                    Ok(None)
190                }
191            })
192            .find_map(|x| x.transpose());
193        if let Some(metadata_file) = metadata_file {
194            let mut entry = metadata_file?;
195            let mut buf = Vec::new();
196            entry.read_to_end(&mut buf)?;
197            Metadata::parse(&buf)
198        } else {
199            Err(Error::MetadataNotFound)
200        }
201    }
202
203    fn parse_zip(path: &Path, metadata_file_suffix: &str) -> Result<Metadata, Error> {
204        let reader = BufReader::new(fs_err::File::open(path)?);
205        let mut archive = ZipArchive::new(reader)?;
206        let metadata_files: Vec<_> = archive
207            .file_names()
208            .filter(|name| name.ends_with(metadata_file_suffix))
209            .map(ToString::to_string)
210            .collect();
211        match metadata_files.as_slice() {
212            [] => Err(Error::MetadataNotFound),
213            [metadata_file] => {
214                let mut buf = Vec::new();
215                archive.by_name(metadata_file)?.read_to_end(&mut buf)?;
216                Metadata::parse(&buf)
217            }
218            [file1, file2]
219                if file1.ends_with(".egg-info/PKG-INFO")
220                    || file2.ends_with(".egg-info/PKG-INFO") =>
221            {
222                let mut buf = Vec::new();
223                archive.by_name(file1)?.read_to_end(&mut buf)?;
224                Metadata::parse(&buf)
225            }
226            _ => {
227                let top_level_files: Vec<_> = metadata_files
228                    .iter()
229                    .filter(|f| {
230                        let path = Path::new(f);
231                        path.components().count() == 2
232                    })
233                    .collect();
234                if top_level_files.len() == 1 {
235                    let mut buf = Vec::new();
236                    archive.by_name(top_level_files[0])?.read_to_end(&mut buf)?;
237                    return Metadata::parse(&buf);
238                }
239                Err(Error::MultipleMetadataFiles(metadata_files))
240            }
241        }
242    }
243}