use std::borrow::Cow;
use std::collections::BTreeMap;
use std::fmt::{self, Debug};
use std::io::{self, Read, Write};
use std::sync::Arc;
use zip::result::{ZipError, ZipResult};
use zip::write::SimpleFileOptions;
use zip::{CompressionMethod, ZipArchive, ZipWriter, read::ZipFile};
pub use self::path::{VendoredPath, VendoredPathBuf};
use crate::file_revision::FileRevision;
mod path;
type Result<T> = io::Result<T>;
#[derive(Clone)]
pub struct VendoredFileSystem {
inner: Arc<VendoredZipArchive>,
}
impl VendoredFileSystem {
pub fn new_static(raw_bytes: &'static [u8]) -> Result<Self> {
Self::new_impl(ArchiveData::Static(raw_bytes))
}
pub fn new(raw_bytes: Vec<u8>) -> Result<Self> {
Self::new_impl(ArchiveData::Owned(raw_bytes.into()))
}
fn new_impl(data: ArchiveData) -> Result<Self> {
Ok(Self {
inner: Arc::new(VendoredZipArchive::new(data)?),
})
}
pub fn exists(&self, path: impl AsRef<VendoredPath>) -> bool {
fn exists(fs: &VendoredFileSystem, path: &VendoredPath) -> bool {
let normalized = NormalizedVendoredPath::from(path);
let archive = &fs.inner;
archive.index_for_path(&normalized).is_some()
|| archive
.index_for_path(&normalized.with_trailing_slash())
.is_some()
}
exists(self, path.as_ref())
}
pub fn metadata(&self, path: impl AsRef<VendoredPath>) -> Result<Metadata> {
fn metadata(fs: &VendoredFileSystem, path: &VendoredPath) -> Result<Metadata> {
let normalized = NormalizedVendoredPath::from(path);
let mut archive = fs.archive_reader();
if let Ok(metadata) = archive.metadata_for_path(&normalized) {
return Ok(metadata);
}
archive.metadata_for_path(&normalized.with_trailing_slash())
}
metadata(self, path.as_ref())
}
pub fn is_directory(&self, path: impl AsRef<VendoredPath>) -> bool {
self.metadata(path)
.is_ok_and(|metadata| metadata.kind().is_directory())
}
pub fn is_file(&self, path: impl AsRef<VendoredPath>) -> bool {
self.metadata(path)
.is_ok_and(|metadata| metadata.kind().is_file())
}
pub fn read_to_string(&self, path: impl AsRef<VendoredPath>) -> Result<String> {
fn read_to_string(fs: &VendoredFileSystem, path: &VendoredPath) -> Result<String> {
let mut archive = fs.archive_reader();
let mut zip_file = archive.lookup_path(&NormalizedVendoredPath::from(path))?;
let mut buffer = String::with_capacity(
usize::try_from(zip_file.size())
.unwrap_or(usize::MAX)
.min(10_000_000),
);
zip_file.read_to_string(&mut buffer)?;
Ok(buffer)
}
read_to_string(self, path.as_ref())
}
pub fn read_directory(
&self,
dir: impl AsRef<VendoredPath>,
) -> impl Iterator<Item = DirectoryEntry> + '_ {
let directory_prefix = NormalizedVendoredPath::from(dir.as_ref())
.with_trailing_slash()
.0
.into_owned();
self.inner.0.file_names().filter_map(move |name| {
let without_dir_prefix = name.strip_prefix(&directory_prefix)?;
if without_dir_prefix.is_empty() {
return None;
}
let file_type = FileType::from_zip_file_name(without_dir_prefix);
let slash_count = without_dir_prefix.matches('/').count();
match file_type {
FileType::File if slash_count > 0 => return None,
FileType::Directory if slash_count > 1 => return None,
_ => {}
}
Some(DirectoryEntry {
path: VendoredPathBuf::from(name),
file_type,
})
})
}
fn archive_reader(&self) -> VendoredZipArchive {
self.inner.as_ref().clone()
}
}
impl fmt::Debug for VendoredFileSystem {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
if f.alternate() {
let mut archive = self.archive_reader();
let mut paths: Vec<String> = archive.0.file_names().map(String::from).collect();
paths.sort();
let debug_info: BTreeMap<String, ZipFileDebugInfo> = paths
.iter()
.map(|path| {
(
path.to_owned(),
ZipFileDebugInfo::from(archive.0.by_name(path).unwrap()),
)
})
.collect();
f.debug_struct("VendoredFileSystem")
.field("paths", &paths)
.field("data_by_path", &debug_info)
.finish()
} else {
write!(f, "VendoredFileSystem(<{} paths>)", self.inner.len())
}
}
}
impl Default for VendoredFileSystem {
fn default() -> Self {
let mut bytes: Vec<u8> = Vec::new();
let mut cursor = io::Cursor::new(&mut bytes);
{
let writer = ZipWriter::new(&mut cursor);
writer.finish().unwrap();
}
VendoredFileSystem::new(bytes).unwrap()
}
}
#[expect(unused)]
#[derive(Debug)]
struct ZipFileDebugInfo {
crc32_hash: u32,
compressed_size: u64,
uncompressed_size: u64,
kind: FileType,
}
impl<'a, R: Read> From<ZipFile<'a, R>> for ZipFileDebugInfo {
fn from(value: ZipFile<'a, R>) -> Self {
Self {
crc32_hash: value.crc32(),
compressed_size: value.compressed_size(),
uncompressed_size: value.size(),
kind: if value.is_dir() {
FileType::Directory
} else {
FileType::File
},
}
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum FileType {
File,
Directory,
}
impl FileType {
fn from_zip_file_name(name: &str) -> FileType {
if name.ends_with('/') {
FileType::Directory
} else {
FileType::File
}
}
pub const fn is_file(self) -> bool {
matches!(self, Self::File)
}
pub const fn is_directory(self) -> bool {
matches!(self, Self::Directory)
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Metadata {
kind: FileType,
revision: FileRevision,
}
impl Metadata {
fn from_zip_file<R: Read>(zip_file: ZipFile<'_, R>) -> Self {
let kind = if zip_file.is_dir() {
FileType::Directory
} else {
FileType::File
};
Self {
kind,
revision: FileRevision::new(u128::from(zip_file.crc32())),
}
}
pub fn kind(&self) -> FileType {
self.kind
}
pub fn revision(&self) -> FileRevision {
self.revision
}
}
#[derive(Debug, PartialEq, Eq)]
pub struct DirectoryEntry {
path: VendoredPathBuf,
file_type: FileType,
}
impl DirectoryEntry {
pub fn new(path: VendoredPathBuf, file_type: FileType) -> Self {
Self { path, file_type }
}
pub fn into_path(self) -> VendoredPathBuf {
self.path
}
pub fn path(&self) -> &VendoredPath {
&self.path
}
pub fn file_type(&self) -> FileType {
self.file_type
}
}
#[derive(Clone, Debug)]
enum ArchiveData {
Static(&'static [u8]),
Owned(Arc<[u8]>),
}
impl AsRef<[u8]> for ArchiveData {
fn as_ref(&self) -> &[u8] {
match self {
Self::Static(data) => data,
Self::Owned(data) => data,
}
}
}
#[derive(Clone, Debug)]
struct VendoredZipArchive(ZipArchive<io::Cursor<ArchiveData>>);
impl VendoredZipArchive {
fn new(data: ArchiveData) -> Result<Self> {
Ok(Self(ZipArchive::new(io::Cursor::new(data))?))
}
fn index_for_path(&self, path: &NormalizedVendoredPath) -> Option<usize> {
self.0.index_for_name(path.as_str())
}
fn lookup_path(
&mut self,
path: &NormalizedVendoredPath,
) -> Result<ZipFile<'_, io::Cursor<ArchiveData>>> {
Ok(self.0.by_name(path.as_str())?)
}
fn metadata_for_path(&mut self, path: &NormalizedVendoredPath) -> Result<Metadata> {
let index = self.index_for_path(path).ok_or(ZipError::FileNotFound)?;
Ok(Metadata::from_zip_file(self.0.by_index_raw(index)?))
}
fn len(&self) -> usize {
self.0.len()
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
struct NormalizedVendoredPath<'a>(Cow<'a, str>);
impl NormalizedVendoredPath<'_> {
fn with_trailing_slash(self) -> Self {
debug_assert!(!self.0.ends_with('/'));
let mut data = self.0.into_owned();
data.push('/');
Self(Cow::Owned(data))
}
fn as_str(&self) -> &str {
&self.0
}
}
impl<'a> From<&'a VendoredPath> for NormalizedVendoredPath<'a> {
fn from(path: &'a VendoredPath) -> Self {
fn normalize_unnormalized_path(path: &VendoredPath) -> String {
let mut normalized_parts = Vec::new();
for component in path.components() {
match component {
camino::Utf8Component::Normal(part) => normalized_parts.push(part),
camino::Utf8Component::CurDir => continue,
camino::Utf8Component::ParentDir => {
normalized_parts.pop();
}
unsupported => {
panic!("Unsupported component in a vendored path: {unsupported}")
}
}
}
normalized_parts.join("/")
}
let path_str = path.as_str();
if std::path::MAIN_SEPARATOR == '\\' && path_str.contains('\\') {
NormalizedVendoredPath(Cow::Owned(normalize_unnormalized_path(path)))
} else if !path
.components()
.all(|component| matches!(component, camino::Utf8Component::Normal(_)))
{
NormalizedVendoredPath(Cow::Owned(normalize_unnormalized_path(path)))
} else {
NormalizedVendoredPath(Cow::Borrowed(path_str.trim_end_matches('/')))
}
}
}
pub struct VendoredFileSystemBuilder {
writer: ZipWriter<io::Cursor<Vec<u8>>>,
compression_method: CompressionMethod,
}
impl VendoredFileSystemBuilder {
pub fn new(compression_method: CompressionMethod) -> Self {
let buffer = io::Cursor::new(Vec::new());
Self {
writer: ZipWriter::new(buffer),
compression_method,
}
}
pub fn add_file(
&mut self,
path: impl AsRef<VendoredPath>,
content: &str,
) -> std::io::Result<()> {
self.writer
.start_file(path.as_ref().as_str(), self.options())?;
self.writer.write_all(content.as_bytes())
}
pub fn add_directory(&mut self, path: impl AsRef<VendoredPath>) -> ZipResult<()> {
self.writer
.add_directory(path.as_ref().as_str(), self.options())
}
pub fn finish(self) -> Result<VendoredFileSystem> {
let buffer = self.writer.finish()?;
VendoredFileSystem::new(buffer.into_inner())
}
fn options(&self) -> SimpleFileOptions {
SimpleFileOptions::default()
.compression_method(self.compression_method)
.unix_permissions(0o644)
}
}
#[cfg(test)]
pub(crate) mod tests {
use insta::assert_snapshot;
use super::*;
const FUNCTOOLS_CONTENTS: &str = "def update_wrapper(): ...";
const ASYNCIO_TASKS_CONTENTS: &str = "class Task: ...";
fn mock_typeshed() -> VendoredFileSystem {
let mut builder = VendoredFileSystemBuilder::new(CompressionMethod::Stored);
builder.add_directory("stdlib/").unwrap();
builder
.add_file("stdlib/functools.pyi", FUNCTOOLS_CONTENTS)
.unwrap();
builder.add_directory("stdlib/asyncio/").unwrap();
builder
.add_file("stdlib/asyncio/tasks.pyi", ASYNCIO_TASKS_CONTENTS)
.unwrap();
builder.finish().unwrap()
}
#[test]
fn filesystem_debug_implementation() {
assert_snapshot!(
format!("{:?}", mock_typeshed()),
@"VendoredFileSystem(<4 paths>)"
);
}
#[test]
fn filesystem_debug_implementation_alternate() {
assert_snapshot!(format!("{:#?}", mock_typeshed()), @r#"
VendoredFileSystem {
paths: [
"stdlib/",
"stdlib/asyncio/",
"stdlib/asyncio/tasks.pyi",
"stdlib/functools.pyi",
],
data_by_path: {
"stdlib/": ZipFileDebugInfo {
crc32_hash: 0,
compressed_size: 0,
uncompressed_size: 0,
kind: Directory,
},
"stdlib/asyncio/": ZipFileDebugInfo {
crc32_hash: 0,
compressed_size: 0,
uncompressed_size: 0,
kind: Directory,
},
"stdlib/asyncio/tasks.pyi": ZipFileDebugInfo {
crc32_hash: 2826547428,
compressed_size: 15,
uncompressed_size: 15,
kind: File,
},
"stdlib/functools.pyi": ZipFileDebugInfo {
crc32_hash: 1099005079,
compressed_size: 25,
uncompressed_size: 25,
kind: File,
},
},
}
"#);
}
fn test_directory(dirname: &str) {
let mock_typeshed = mock_typeshed();
let path = VendoredPath::new(dirname);
assert!(mock_typeshed.exists(path));
assert!(mock_typeshed.read_to_string(path).is_err());
let metadata = mock_typeshed.metadata(path).unwrap();
assert!(metadata.kind().is_directory());
}
#[test]
fn stdlib_dir_no_trailing_slash() {
test_directory("stdlib")
}
#[test]
fn stdlib_dir_trailing_slash() {
test_directory("stdlib/")
}
#[test]
fn asyncio_dir_no_trailing_slash() {
test_directory("stdlib/asyncio")
}
#[test]
fn asyncio_dir_trailing_slash() {
test_directory("stdlib/asyncio/")
}
#[test]
fn stdlib_dir_parent_components() {
test_directory("stdlib/asyncio/../../stdlib")
}
#[test]
fn asyncio_dir_odd_components() {
test_directory("./stdlib/asyncio/../asyncio/")
}
fn readdir_snapshot(fs: &VendoredFileSystem, path: &str) -> String {
let mut paths = fs
.read_directory(VendoredPath::new(path))
.map(|entry| entry.path().to_string())
.collect::<Vec<String>>();
paths.sort();
paths.join("\n")
}
#[test]
fn read_directory_stdlib() {
let mock_typeshed = mock_typeshed();
assert_snapshot!(readdir_snapshot(&mock_typeshed, "stdlib"), @"
vendored://stdlib/asyncio/
vendored://stdlib/functools.pyi
");
assert_snapshot!(readdir_snapshot(&mock_typeshed, "stdlib/"), @"
vendored://stdlib/asyncio/
vendored://stdlib/functools.pyi
");
assert_snapshot!(readdir_snapshot(&mock_typeshed, "./stdlib"), @"
vendored://stdlib/asyncio/
vendored://stdlib/functools.pyi
");
assert_snapshot!(readdir_snapshot(&mock_typeshed, "./stdlib/"), @"
vendored://stdlib/asyncio/
vendored://stdlib/functools.pyi
");
}
#[test]
fn read_directory_asyncio() {
let mock_typeshed = mock_typeshed();
assert_snapshot!(
readdir_snapshot(&mock_typeshed, "stdlib/asyncio"),
@"vendored://stdlib/asyncio/tasks.pyi",
);
assert_snapshot!(
readdir_snapshot(&mock_typeshed, "./stdlib/asyncio"),
@"vendored://stdlib/asyncio/tasks.pyi",
);
assert_snapshot!(
readdir_snapshot(&mock_typeshed, "stdlib/asyncio/"),
@"vendored://stdlib/asyncio/tasks.pyi",
);
assert_snapshot!(
readdir_snapshot(&mock_typeshed, "./stdlib/asyncio/"),
@"vendored://stdlib/asyncio/tasks.pyi",
);
}
fn test_nonexistent_path(path: &str) {
let mock_typeshed = mock_typeshed();
let path = VendoredPath::new(path);
assert!(!mock_typeshed.exists(path));
assert!(mock_typeshed.metadata(path).is_err());
assert!(
mock_typeshed
.read_to_string(path)
.is_err_and(|err| err.to_string().contains("file not found"))
);
}
#[test]
fn simple_nonexistent_path() {
test_nonexistent_path("foo")
}
#[test]
fn nonexistent_path_with_extension() {
test_nonexistent_path("foo.pyi")
}
#[test]
fn nonexistent_path_with_trailing_slash() {
test_nonexistent_path("foo/")
}
#[test]
fn nonexistent_path_with_fancy_components() {
test_nonexistent_path("./foo/../../../foo")
}
fn test_file(mock_typeshed: &VendoredFileSystem, path: &VendoredPath) {
assert!(mock_typeshed.exists(path));
let metadata = mock_typeshed.metadata(path).unwrap();
assert!(metadata.kind().is_file());
}
#[test]
fn functools_file_contents() {
let mock_typeshed = mock_typeshed();
let path = VendoredPath::new("stdlib/functools.pyi");
test_file(&mock_typeshed, path);
let functools_stub = mock_typeshed.read_to_string(path).unwrap();
assert_eq!(functools_stub.as_str(), FUNCTOOLS_CONTENTS);
let functools_stub_again = mock_typeshed.read_to_string(path).unwrap();
assert_eq!(functools_stub_again.as_str(), FUNCTOOLS_CONTENTS);
}
#[test]
fn functools_file_other_path() {
test_file(
&mock_typeshed(),
VendoredPath::new("stdlib/../stdlib/../stdlib/functools.pyi"),
)
}
#[test]
fn asyncio_file_contents() {
let mock_typeshed = mock_typeshed();
let path = VendoredPath::new("stdlib/asyncio/tasks.pyi");
test_file(&mock_typeshed, path);
let asyncio_stub = mock_typeshed.read_to_string(path).unwrap();
assert_eq!(asyncio_stub.as_str(), ASYNCIO_TASKS_CONTENTS);
}
#[test]
fn asyncio_file_other_path() {
test_file(
&mock_typeshed(),
VendoredPath::new("./stdlib/asyncio/../asyncio/tasks.pyi"),
)
}
}