torsh_package/
resources.rs

1//! Resource handling for packages
2
3use serde::{Deserialize, Serialize};
4use std::collections::HashMap;
5use std::path::Path;
6
7/// Type of resource in the package
8#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
9pub enum ResourceType {
10    /// Model weights and state
11    Model,
12    /// Source code
13    Source,
14    /// Data files
15    Data,
16    /// Configuration files
17    Config,
18    /// Documentation
19    Documentation,
20    /// License files
21    License,
22    /// Binary assets
23    Binary,
24    /// Text assets
25    Text,
26    /// Metadata
27    Metadata,
28}
29
30impl ResourceType {
31    /// Get file extension for resource type
32    pub fn extension(&self) -> &'static str {
33        match self {
34            ResourceType::Model => "model",
35            ResourceType::Source => "rs",
36            ResourceType::Data => "data",
37            ResourceType::Config => "toml",
38            ResourceType::Documentation => "md",
39            ResourceType::License => "txt",
40            ResourceType::Binary => "bin",
41            ResourceType::Text => "txt",
42            ResourceType::Metadata => "json",
43        }
44    }
45
46    /// Get MIME type for resource
47    pub fn mime_type(&self) -> &'static str {
48        match self {
49            ResourceType::Model => "application/octet-stream",
50            ResourceType::Source => "text/x-rust",
51            ResourceType::Data => "application/octet-stream",
52            ResourceType::Config => "text/x-toml",
53            ResourceType::Documentation => "text/markdown",
54            ResourceType::License => "text/plain",
55            ResourceType::Binary => "application/octet-stream",
56            ResourceType::Text => "text/plain",
57            ResourceType::Metadata => "application/json",
58        }
59    }
60
61    /// Determine resource type from file extension
62    pub fn from_extension(ext: &str) -> Self {
63        match ext.to_lowercase().as_str() {
64            "model" | "pth" | "pt" | "torsh" => ResourceType::Model,
65            "rs" => ResourceType::Source,
66            "json" => ResourceType::Metadata,
67            "toml" | "yaml" | "yml" => ResourceType::Config,
68            "md" | "rst" => ResourceType::Documentation,
69            "txt" | "license" => ResourceType::License,
70            "bin" | "dat" => ResourceType::Binary,
71            _ => ResourceType::Data,
72        }
73    }
74
75    /// Canonical, stable string tag for this resource type.
76    ///
77    /// Used by [`crate::exporter::PackageExporter`] to record a resource's
78    /// exact type in its `.metadata` side-file (see
79    /// [`RESOURCE_TYPE_METADATA_KEY`]), so [`crate::importer::PackageImporter`]
80    /// can recover it authoritatively on import instead of re-deriving it
81    /// from the archive path prefix or file extension — both of which are
82    /// ambiguous (every type other than Model/Source/Data/Config/Documentation
83    /// shares the same `resources/` archive prefix, and file extensions are
84    /// not a 1:1 mapping back to a single [`ResourceType`]). See
85    /// [`Self::from_tag`] for the inverse.
86    pub fn as_tag(&self) -> &'static str {
87        match self {
88            ResourceType::Model => "model",
89            ResourceType::Source => "source",
90            ResourceType::Data => "data",
91            ResourceType::Config => "config",
92            ResourceType::Documentation => "documentation",
93            ResourceType::License => "license",
94            ResourceType::Binary => "binary",
95            ResourceType::Text => "text",
96            ResourceType::Metadata => "metadata",
97        }
98    }
99
100    /// Parse a tag produced by [`Self::as_tag`].
101    ///
102    /// Returns `None` for an unrecognized tag (e.g. a package written by a
103    /// version of this crate that used a different tag set), so callers can
104    /// fall back to prefix/extension-based classification rather than
105    /// failing the import outright.
106    pub fn from_tag(tag: &str) -> Option<Self> {
107        match tag {
108            "model" => Some(ResourceType::Model),
109            "source" => Some(ResourceType::Source),
110            "data" => Some(ResourceType::Data),
111            "config" => Some(ResourceType::Config),
112            "documentation" => Some(ResourceType::Documentation),
113            "license" => Some(ResourceType::License),
114            "binary" => Some(ResourceType::Binary),
115            "text" => Some(ResourceType::Text),
116            "metadata" => Some(ResourceType::Metadata),
117            _ => None,
118        }
119    }
120}
121
122/// Reserved [`Resource::metadata`] key used to persist a resource's exact
123/// [`ResourceType`] in its `.metadata` archive side-file across an
124/// export/import roundtrip (F240).
125///
126/// Chosen to be extremely unlikely to collide with a caller-supplied
127/// metadata key. [`crate::importer::PackageImporter`] strips this key back
128/// out of the metadata map it exposes on the imported [`Resource`], so it
129/// never leaks into user-visible metadata.
130pub(crate) const RESOURCE_TYPE_METADATA_KEY: &str = "__torsh_resource_type";
131
132/// A resource in the package
133#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
134pub struct Resource {
135    /// Resource name (unique identifier)
136    pub name: String,
137
138    /// Resource type
139    pub resource_type: ResourceType,
140
141    /// Resource data
142    pub data: Vec<u8>,
143
144    /// Additional metadata
145    pub metadata: HashMap<String, String>,
146}
147
148impl Resource {
149    /// Create a new resource
150    pub fn new(name: String, resource_type: ResourceType, data: Vec<u8>) -> Self {
151        Self {
152            name,
153            resource_type,
154            data,
155            metadata: HashMap::new(),
156        }
157    }
158
159    /// Create resource from file
160    pub fn from_file<P: AsRef<Path>>(
161        path: P,
162        resource_type: Option<ResourceType>,
163    ) -> std::io::Result<Self> {
164        let path = path.as_ref();
165        let name = path
166            .file_name()
167            .and_then(|n| n.to_str())
168            .unwrap_or("unnamed")
169            .to_string();
170
171        let resource_type = resource_type.unwrap_or_else(|| {
172            path.extension()
173                .and_then(|e| e.to_str())
174                .map(ResourceType::from_extension)
175                .unwrap_or(ResourceType::Data)
176        });
177
178        let data = std::fs::read(path)?;
179
180        Ok(Self::new(name, resource_type, data))
181    }
182
183    /// Get resource size
184    pub fn size(&self) -> usize {
185        self.data.len()
186    }
187
188    /// Calculate SHA256 hash
189    pub fn sha256(&self) -> String {
190        use sha2::{Digest, Sha256};
191        let mut hasher = Sha256::new();
192        hasher.update(&self.data);
193        hex::encode(hasher.finalize())
194    }
195
196    /// Add metadata
197    pub fn add_metadata(&mut self, key: String, value: String) {
198        self.metadata.insert(key, value);
199    }
200
201    /// Get metadata value
202    pub fn get_metadata(&self, key: &str) -> Option<&str> {
203        self.metadata.get(key).map(|s| s.as_str())
204    }
205
206    /// Check if resource is compressed
207    pub fn is_compressed(&self) -> bool {
208        self.metadata.contains_key("compression")
209    }
210
211    /// Get compression method
212    pub fn compression_method(&self) -> Option<&str> {
213        self.get_metadata("compression")
214    }
215}
216
217/// Resource collection for managing multiple resources
218pub struct ResourceCollection {
219    resources: HashMap<String, Resource>,
220}
221
222impl ResourceCollection {
223    /// Create a new resource collection
224    pub fn new() -> Self {
225        Self {
226            resources: HashMap::new(),
227        }
228    }
229
230    /// Add a resource
231    pub fn add(&mut self, resource: Resource) -> Result<(), String> {
232        if self.resources.contains_key(&resource.name) {
233            return Err(format!("Resource '{}' already exists", resource.name));
234        }
235
236        self.resources.insert(resource.name.clone(), resource);
237        Ok(())
238    }
239
240    /// Get a resource by name
241    pub fn get(&self, name: &str) -> Option<&Resource> {
242        self.resources.get(name)
243    }
244
245    /// Get mutable resource by name
246    pub fn get_mut(&mut self, name: &str) -> Option<&mut Resource> {
247        self.resources.get_mut(name)
248    }
249
250    /// Remove a resource
251    pub fn remove(&mut self, name: &str) -> Option<Resource> {
252        self.resources.remove(name)
253    }
254
255    /// List all resource names
256    pub fn list(&self) -> Vec<&str> {
257        self.resources.keys().map(|s| s.as_str()).collect()
258    }
259
260    /// Get resources by type
261    pub fn by_type(&self, resource_type: ResourceType) -> Vec<&Resource> {
262        self.resources
263            .values()
264            .filter(|r| r.resource_type == resource_type)
265            .collect()
266    }
267
268    /// Get total size of all resources
269    pub fn total_size(&self) -> usize {
270        self.resources.values().map(|r| r.size()).sum()
271    }
272
273    /// Clear all resources
274    pub fn clear(&mut self) {
275        self.resources.clear();
276    }
277
278    /// Get number of resources
279    pub fn len(&self) -> usize {
280        self.resources.len()
281    }
282
283    /// Check if empty
284    pub fn is_empty(&self) -> bool {
285        self.resources.is_empty()
286    }
287}
288
289impl Default for ResourceCollection {
290    fn default() -> Self {
291        Self::new()
292    }
293}
294
295/// Resource filter for selective operations
296pub struct ResourceFilter {
297    types: Option<Vec<ResourceType>>,
298    name_pattern: Option<String>,
299    min_size: Option<usize>,
300    max_size: Option<usize>,
301}
302
303impl ResourceFilter {
304    /// Create a new filter
305    pub fn new() -> Self {
306        Self {
307            types: None,
308            name_pattern: None,
309            min_size: None,
310            max_size: None,
311        }
312    }
313
314    /// Filter by resource types
315    pub fn with_types(mut self, types: Vec<ResourceType>) -> Self {
316        self.types = Some(types);
317        self
318    }
319
320    /// Filter by name pattern
321    pub fn with_name_pattern(mut self, pattern: String) -> Self {
322        self.name_pattern = Some(pattern);
323        self
324    }
325
326    /// Filter by size range
327    pub fn with_size_range(mut self, min: Option<usize>, max: Option<usize>) -> Self {
328        self.min_size = min;
329        self.max_size = max;
330        self
331    }
332
333    /// Check if resource matches filter
334    pub fn matches(&self, resource: &Resource) -> bool {
335        // Check type filter
336        if let Some(types) = &self.types {
337            if !types.contains(&resource.resource_type) {
338                return false;
339            }
340        }
341
342        // Check name pattern
343        if let Some(pattern) = &self.name_pattern {
344            if !resource.name.contains(pattern) {
345                return false;
346            }
347        }
348
349        // Check size range
350        let size = resource.size();
351        if let Some(min) = self.min_size {
352            if size < min {
353                return false;
354            }
355        }
356        if let Some(max) = self.max_size {
357            if size > max {
358                return false;
359            }
360        }
361
362        true
363    }
364}
365
366impl Default for ResourceFilter {
367    fn default() -> Self {
368        Self::new()
369    }
370}
371
372#[cfg(test)]
373mod tests {
374    use super::*;
375
376    #[test]
377    fn test_resource_type() {
378        assert_eq!(ResourceType::from_extension("rs"), ResourceType::Source);
379        assert_eq!(ResourceType::from_extension("model"), ResourceType::Model);
380        assert_eq!(ResourceType::from_extension("json"), ResourceType::Metadata);
381        assert_eq!(ResourceType::from_extension("unknown"), ResourceType::Data);
382    }
383
384    #[test]
385    fn test_resource_creation() {
386        let resource = Resource::new(
387            "test.model".to_string(),
388            ResourceType::Model,
389            vec![1, 2, 3, 4],
390        );
391
392        assert_eq!(resource.name, "test.model");
393        assert_eq!(resource.resource_type, ResourceType::Model);
394        assert_eq!(resource.size(), 4);
395    }
396
397    #[test]
398    fn test_resource_collection() {
399        let mut collection = ResourceCollection::new();
400
401        let resource1 = Resource::new("res1".to_string(), ResourceType::Model, vec![1, 2, 3]);
402        let resource2 = Resource::new("res2".to_string(), ResourceType::Data, vec![4, 5, 6, 7]);
403
404        collection.add(resource1).unwrap();
405        collection.add(resource2).unwrap();
406
407        assert_eq!(collection.len(), 2);
408        assert_eq!(collection.total_size(), 7);
409
410        let models = collection.by_type(ResourceType::Model);
411        assert_eq!(models.len(), 1);
412    }
413
414    #[test]
415    fn test_resource_filter() {
416        let resource = Resource::new("test.model".to_string(), ResourceType::Model, vec![0; 100]);
417
418        let filter = ResourceFilter::new()
419            .with_types(vec![ResourceType::Model, ResourceType::Data])
420            .with_size_range(Some(50), Some(200));
421
422        assert!(filter.matches(&resource));
423
424        let filter2 = ResourceFilter::new()
425            .with_types(vec![ResourceType::Source])
426            .with_size_range(Some(50), Some(200));
427
428        assert!(!filter2.matches(&resource));
429    }
430}
torsh_package/resources.rs

torsh_package/
resources.rs
