burn_central_core/bundle/memory/
sources.rs

1use std::io::Read;
2
3use crate::bundle::{BundleSink, normalize_bundle_path};
4
5/// A builder for creating bundles with multiple files
6#[derive(Default, Clone)]
7pub struct InMemoryBundleSources {
8    files: Vec<PendingFile>,
9}
10
11impl InMemoryBundleSources {
12    /// Create a new empty bundle sources
13    pub fn new() -> Self {
14        Self::default()
15    }
16
17    /// Add raw bytes as a file within the bundle at `dest_path`.
18    pub fn add_bytes(mut self, bytes: Vec<u8>, dest_path: impl AsRef<str>) -> Self {
19        self.files.push(PendingFile {
20            dest_path: normalize_bundle_path(dest_path.as_ref()),
21            source: bytes,
22        });
23        self
24    }
25
26    /// Add a file from a reader
27    pub fn add_file<R: Read>(
28        self,
29        mut reader: R,
30        dest_path: impl AsRef<str>,
31    ) -> Result<Self, std::io::Error> {
32        let mut bytes = Vec::new();
33        reader.read_to_end(&mut bytes)?;
34        Ok(self.add_bytes(bytes, dest_path))
35    }
36
37    /// Get the files in this bundle sources
38    pub fn files(&self) -> &Vec<PendingFile> {
39        &self.files
40    }
41
42    /// Convert into the files vector
43    pub fn into_files(self) -> Vec<PendingFile> {
44        self.files
45    }
46
47    /// Check if the bundle is empty
48    pub fn is_empty(&self) -> bool {
49        self.files.is_empty()
50    }
51
52    /// Get the number of files
53    pub fn len(&self) -> usize {
54        self.files.len()
55    }
56}
57
58impl BundleSink for InMemoryBundleSources {
59    fn put_file<R: Read>(&mut self, path: &str, reader: &mut R) -> Result<(), String> {
60        let mut buf = Vec::new();
61        reader
62            .read_to_end(&mut buf)
63            .map_err(|e| format!("Failed to read from source: {}", e))?;
64        *self = self.clone().add_bytes(buf, path);
65        Ok(())
66    }
67}
68
69/// A file that is pending to be added to a bundle
70#[derive(Clone)]
71pub struct PendingFile {
72    pub dest_path: String, // path within the bundle (use forward slashes)
73    pub source: Vec<u8>,
74}
75
76impl PendingFile {
77    /// Get the destination path of this file
78    pub fn dest_path(&self) -> &str {
79        &self.dest_path
80    }
81
82    /// Get the source bytes of this file
83    pub fn source(&self) -> &[u8] {
84        &self.source
85    }
86
87    /// Get the size of this file
88    pub fn size(&self) -> usize {
89        self.source.len()
90    }
91}
burn_central_core/bundle/memory/sources.rs

burn_central_core/bundle/memory/
sources.rs
