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
76/// A resource in the package
77#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
78pub struct Resource {
79    /// Resource name (unique identifier)
80    pub name: String,
81
82    /// Resource type
83    pub resource_type: ResourceType,
84
85    /// Resource data
86    pub data: Vec<u8>,
87
88    /// Additional metadata
89    pub metadata: HashMap<String, String>,
90}
91
92impl Resource {
93    /// Create a new resource
94    pub fn new(name: String, resource_type: ResourceType, data: Vec<u8>) -> Self {
95        Self {
96            name,
97            resource_type,
98            data,
99            metadata: HashMap::new(),
100        }
101    }
102
103    /// Create resource from file
104    pub fn from_file<P: AsRef<Path>>(
105        path: P,
106        resource_type: Option<ResourceType>,
107    ) -> std::io::Result<Self> {
108        let path = path.as_ref();
109        let name = path
110            .file_name()
111            .and_then(|n| n.to_str())
112            .unwrap_or("unnamed")
113            .to_string();
114
115        let resource_type = resource_type.unwrap_or_else(|| {
116            path.extension()
117                .and_then(|e| e.to_str())
118                .map(ResourceType::from_extension)
119                .unwrap_or(ResourceType::Data)
120        });
121
122        let data = std::fs::read(path)?;
123
124        Ok(Self::new(name, resource_type, data))
125    }
126
127    /// Get resource size
128    pub fn size(&self) -> usize {
129        self.data.len()
130    }
131
132    /// Calculate SHA256 hash
133    pub fn sha256(&self) -> String {
134        use sha2::{Digest, Sha256};
135        let mut hasher = Sha256::new();
136        hasher.update(&self.data);
137        hex::encode(hasher.finalize())
138    }
139
140    /// Add metadata
141    pub fn add_metadata(&mut self, key: String, value: String) {
142        self.metadata.insert(key, value);
143    }
144
145    /// Get metadata value
146    pub fn get_metadata(&self, key: &str) -> Option<&str> {
147        self.metadata.get(key).map(|s| s.as_str())
148    }
149
150    /// Check if resource is compressed
151    pub fn is_compressed(&self) -> bool {
152        self.metadata.contains_key("compression")
153    }
154
155    /// Get compression method
156    pub fn compression_method(&self) -> Option<&str> {
157        self.get_metadata("compression")
158    }
159}
160
161/// Resource collection for managing multiple resources
162pub struct ResourceCollection {
163    resources: HashMap<String, Resource>,
164}
165
166impl ResourceCollection {
167    /// Create a new resource collection
168    pub fn new() -> Self {
169        Self {
170            resources: HashMap::new(),
171        }
172    }
173
174    /// Add a resource
175    pub fn add(&mut self, resource: Resource) -> Result<(), String> {
176        if self.resources.contains_key(&resource.name) {
177            return Err(format!("Resource '{}' already exists", resource.name));
178        }
179
180        self.resources.insert(resource.name.clone(), resource);
181        Ok(())
182    }
183
184    /// Get a resource by name
185    pub fn get(&self, name: &str) -> Option<&Resource> {
186        self.resources.get(name)
187    }
188
189    /// Get mutable resource by name
190    pub fn get_mut(&mut self, name: &str) -> Option<&mut Resource> {
191        self.resources.get_mut(name)
192    }
193
194    /// Remove a resource
195    pub fn remove(&mut self, name: &str) -> Option<Resource> {
196        self.resources.remove(name)
197    }
198
199    /// List all resource names
200    pub fn list(&self) -> Vec<&str> {
201        self.resources.keys().map(|s| s.as_str()).collect()
202    }
203
204    /// Get resources by type
205    pub fn by_type(&self, resource_type: ResourceType) -> Vec<&Resource> {
206        self.resources
207            .values()
208            .filter(|r| r.resource_type == resource_type)
209            .collect()
210    }
211
212    /// Get total size of all resources
213    pub fn total_size(&self) -> usize {
214        self.resources.values().map(|r| r.size()).sum()
215    }
216
217    /// Clear all resources
218    pub fn clear(&mut self) {
219        self.resources.clear();
220    }
221
222    /// Get number of resources
223    pub fn len(&self) -> usize {
224        self.resources.len()
225    }
226
227    /// Check if empty
228    pub fn is_empty(&self) -> bool {
229        self.resources.is_empty()
230    }
231}
232
233impl Default for ResourceCollection {
234    fn default() -> Self {
235        Self::new()
236    }
237}
238
239/// Resource filter for selective operations
240pub struct ResourceFilter {
241    types: Option<Vec<ResourceType>>,
242    name_pattern: Option<String>,
243    min_size: Option<usize>,
244    max_size: Option<usize>,
245}
246
247impl ResourceFilter {
248    /// Create a new filter
249    pub fn new() -> Self {
250        Self {
251            types: None,
252            name_pattern: None,
253            min_size: None,
254            max_size: None,
255        }
256    }
257
258    /// Filter by resource types
259    pub fn with_types(mut self, types: Vec<ResourceType>) -> Self {
260        self.types = Some(types);
261        self
262    }
263
264    /// Filter by name pattern
265    pub fn with_name_pattern(mut self, pattern: String) -> Self {
266        self.name_pattern = Some(pattern);
267        self
268    }
269
270    /// Filter by size range
271    pub fn with_size_range(mut self, min: Option<usize>, max: Option<usize>) -> Self {
272        self.min_size = min;
273        self.max_size = max;
274        self
275    }
276
277    /// Check if resource matches filter
278    pub fn matches(&self, resource: &Resource) -> bool {
279        // Check type filter
280        if let Some(types) = &self.types {
281            if !types.contains(&resource.resource_type) {
282                return false;
283            }
284        }
285
286        // Check name pattern
287        if let Some(pattern) = &self.name_pattern {
288            if !resource.name.contains(pattern) {
289                return false;
290            }
291        }
292
293        // Check size range
294        let size = resource.size();
295        if let Some(min) = self.min_size {
296            if size < min {
297                return false;
298            }
299        }
300        if let Some(max) = self.max_size {
301            if size > max {
302                return false;
303            }
304        }
305
306        true
307    }
308}
309
310impl Default for ResourceFilter {
311    fn default() -> Self {
312        Self::new()
313    }
314}
315
316#[cfg(test)]
317mod tests {
318    use super::*;
319
320    #[test]
321    fn test_resource_type() {
322        assert_eq!(ResourceType::from_extension("rs"), ResourceType::Source);
323        assert_eq!(ResourceType::from_extension("model"), ResourceType::Model);
324        assert_eq!(ResourceType::from_extension("json"), ResourceType::Metadata);
325        assert_eq!(ResourceType::from_extension("unknown"), ResourceType::Data);
326    }
327
328    #[test]
329    fn test_resource_creation() {
330        let resource = Resource::new(
331            "test.model".to_string(),
332            ResourceType::Model,
333            vec![1, 2, 3, 4],
334        );
335
336        assert_eq!(resource.name, "test.model");
337        assert_eq!(resource.resource_type, ResourceType::Model);
338        assert_eq!(resource.size(), 4);
339    }
340
341    #[test]
342    fn test_resource_collection() {
343        let mut collection = ResourceCollection::new();
344
345        let resource1 = Resource::new("res1".to_string(), ResourceType::Model, vec![1, 2, 3]);
346        let resource2 = Resource::new("res2".to_string(), ResourceType::Data, vec![4, 5, 6, 7]);
347
348        collection.add(resource1).unwrap();
349        collection.add(resource2).unwrap();
350
351        assert_eq!(collection.len(), 2);
352        assert_eq!(collection.total_size(), 7);
353
354        let models = collection.by_type(ResourceType::Model);
355        assert_eq!(models.len(), 1);
356    }
357
358    #[test]
359    fn test_resource_filter() {
360        let resource = Resource::new("test.model".to_string(), ResourceType::Model, vec![0; 100]);
361
362        let filter = ResourceFilter::new()
363            .with_types(vec![ResourceType::Model, ResourceType::Data])
364            .with_size_range(Some(50), Some(200));
365
366        assert!(filter.matches(&resource));
367
368        let filter2 = ResourceFilter::new()
369            .with_types(vec![ResourceType::Source])
370            .with_size_range(Some(50), Some(200));
371
372        assert!(!filter2.matches(&resource));
373    }
374}
torsh_package/resources.rs

torsh_package/
resources.rs
