dataforge/customization/
mod.rs

1//! 用户自定义模块
2//! 
3//! 提供用户自定义数据生成规则和格式的支持,满足个性化需求
4
5use std::collections::HashMap;
6use serde_json::Value;
7use crate::error::{DataForgeError, Result};
8
9pub mod api;
10pub mod plugins;
11
12pub use api::*;
13pub use plugins::*;
14
15/// 自定义生成器特征
16pub trait CustomGenerator: Send + Sync {
17    /// 生成器名称
18    fn name(&self) -> &str;
19    
20    /// 生成器描述
21    fn description(&self) -> &str;
22    
23    /// 生成数据
24    fn generate(&self) -> Result<Value>;
25    
26    /// 生成带参数的数据
27    fn generate_with_params(&self, params: &HashMap<String, Value>) -> Result<Value> {
28        // 默认实现忽略参数
29        let _ = params;
30        self.generate()
31    }
32    
33    /// 验证参数
34    fn validate_params(&self, params: &HashMap<String, Value>) -> Result<()> {
35        // 默认实现不验证参数
36        let _ = params;
37        Ok(())
38    }
39    
40    /// 获取参数模式
41    fn param_schema(&self) -> HashMap<String, ParamType> {
42        HashMap::new()
43    }
44}
45
46/// 参数类型
47#[derive(Debug, Clone, PartialEq)]
48pub enum ParamType {
49    String { max_length: Option<usize> },
50    Integer { min: Option<i64>, max: Option<i64> },
51    Float { min: Option<f64>, max: Option<f64> },
52    Boolean,
53    Array(Box<ParamType>),
54    Object(HashMap<String, ParamType>),
55}
56
57/// 自定义生成器注册表
58pub struct CustomGeneratorRegistry {
59    generators: HashMap<String, Box<dyn CustomGenerator>>,
60}
61
62impl CustomGeneratorRegistry {
63    /// 创建新的注册表
64    pub fn new() -> Self {
65        Self {
66            generators: HashMap::new(),
67        }
68    }
69
70    /// 注册生成器
71    pub fn register<G>(&mut self, generator: G) -> Result<()>
72    where
73        G: CustomGenerator + 'static,
74    {
75        let name = generator.name().to_string();
76        
77        if self.generators.contains_key(&name) {
78            return Err(DataForgeError::validation(&format!(
79                "Generator '{}' is already registered", name
80            )));
81        }
82        
83        self.generators.insert(name, Box::new(generator));
84        Ok(())
85    }
86
87    /// 获取生成器
88    pub fn get(&self, name: &str) -> Option<&dyn CustomGenerator> {
89        self.generators.get(name).map(|g| g.as_ref())
90    }
91
92    /// 移除生成器
93    pub fn unregister(&mut self, name: &str) -> bool {
94        self.generators.remove(name).is_some()
95    }
96
97    /// 列出所有生成器
98    pub fn list_generators(&self) -> Vec<GeneratorInfo> {
99        self.generators.iter().map(|(name, generator)| {
100            GeneratorInfo {
101                name: name.clone(),
102                description: generator.description().to_string(),
103                param_schema: generator.param_schema(),
104            }
105        }).collect()
106    }
107
108    /// 生成数据
109    pub fn generate(&self, name: &str, params: Option<&HashMap<String, Value>>) -> Result<Value> {
110        let generator = self.generators.get(name)
111            .ok_or_else(|| DataForgeError::validation(&format!("Generator '{}' not found", name)))?;
112
113        if let Some(params) = params {
114            generator.validate_params(params)?;
115            generator.generate_with_params(params)
116        } else {
117            generator.generate()
118        }
119    }
120
121    /// 获取生成器数量
122    pub fn count(&self) -> usize {
123        self.generators.len()
124    }
125}
126
127impl Default for CustomGeneratorRegistry {
128    fn default() -> Self {
129        Self::new()
130    }
131}
132
133/// 生成器信息
134#[derive(Debug, Clone)]
135pub struct GeneratorInfo {
136    pub name: String,
137    pub description: String,
138    pub param_schema: HashMap<String, ParamType>,
139}
140
141/// 内置自定义生成器示例
142pub struct RandomStringGenerator {
143    name: String,
144}
145
146impl RandomStringGenerator {
147    pub fn new() -> Self {
148        Self {
149            name: "random_string".to_string(),
150        }
151    }
152}
153
154impl CustomGenerator for RandomStringGenerator {
155    fn name(&self) -> &str {
156        &self.name
157    }
158
159    fn description(&self) -> &str {
160        "Generates random strings with configurable length and character set"
161    }
162
163    fn generate(&self) -> Result<Value> {
164        use rand::Rng;
165        let mut rng = rand::thread_rng();
166        let length = rng.gen_range(5..=20);
167        let chars: Vec<char> = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789".chars().collect();
168        
169        let random_string: String = (0..length)
170            .map(|_| chars[rng.gen_range(0..chars.len())])
171            .collect();
172            
173        Ok(Value::String(random_string))
174    }
175
176    fn generate_with_params(&self, params: &HashMap<String, Value>) -> Result<Value> {
177        use rand::Rng;
178        let mut rng = rand::thread_rng();
179        
180        let length = params.get("length")
181            .and_then(|v| v.as_u64())
182            .unwrap_or(10) as usize;
183            
184        let charset = params.get("charset")
185            .and_then(|v| v.as_str())
186            .unwrap_or("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789");
187            
188        let chars: Vec<char> = charset.chars().collect();
189        if chars.is_empty() {
190            return Err(DataForgeError::validation("Character set cannot be empty"));
191        }
192        
193        let random_string: String = (0..length)
194            .map(|_| chars[rng.gen_range(0..chars.len())])
195            .collect();
196            
197        Ok(Value::String(random_string))
198    }
199
200    fn validate_params(&self, params: &HashMap<String, Value>) -> Result<()> {
201        if let Some(length) = params.get("length") {
202            if let Some(len) = length.as_u64() {
203                if len == 0 || len > 1000 {
204                    return Err(DataForgeError::validation("Length must be between 1 and 1000"));
205                }
206            } else {
207                return Err(DataForgeError::validation("Length must be a positive integer"));
208            }
209        }
210        
211        if let Some(charset) = params.get("charset") {
212            if let Some(cs) = charset.as_str() {
213                if cs.is_empty() {
214                    return Err(DataForgeError::validation("Character set cannot be empty"));
215                }
216            } else {
217                return Err(DataForgeError::validation("Character set must be a string"));
218            }
219        }
220        
221        Ok(())
222    }
223
224    fn param_schema(&self) -> HashMap<String, ParamType> {
225        let mut schema = HashMap::new();
226        schema.insert("length".to_string(), ParamType::Integer { min: Some(1), max: Some(1000) });
227        schema.insert("charset".to_string(), ParamType::String { max_length: None });
228        schema
229    }
230}
231
232impl Default for RandomStringGenerator {
233    fn default() -> Self {
234        Self::new()
235    }
236}
237
238#[cfg(test)]
239mod tests {
240    use super::*;
241
242    #[test]
243    fn test_custom_generator_registry() {
244        let mut registry = CustomGeneratorRegistry::new();
245        let generator = RandomStringGenerator::new();
246        
247        assert!(registry.register(generator).is_ok());
248        assert_eq!(registry.count(), 1);
249        
250        let result = registry.generate("random_string", None);
251        assert!(result.is_ok());
252        
253        if let Ok(Value::String(s)) = result {
254            assert!(!s.is_empty());
255        } else {
256            panic!("Expected string value");
257        }
258    }
259
260    #[test]
261    fn test_random_string_generator() {
262        let generator = RandomStringGenerator::new();
263        
264        assert_eq!(generator.name(), "random_string");
265        assert!(!generator.description().is_empty());
266        
267        let result = generator.generate();
268        assert!(result.is_ok());
269        
270        if let Ok(Value::String(s)) = result {
271            assert!(s.len() >= 5 && s.len() <= 20);
272        }
273    }
274
275    #[test]
276    fn test_random_string_generator_with_params() {
277        let generator = RandomStringGenerator::new();
278        let mut params = HashMap::new();
279        params.insert("length".to_string(), Value::Number(serde_json::Number::from(5)));
280        params.insert("charset".to_string(), Value::String("abc".to_string()));
281        
282        assert!(generator.validate_params(&params).is_ok());
283        
284        let result = generator.generate_with_params(&params);
285        assert!(result.is_ok());
286        
287        if let Ok(Value::String(s)) = result {
288            assert_eq!(s.len(), 5);
289            assert!(s.chars().all(|c| "abc".contains(c)));
290        }
291    }
292
293    #[test]
294    fn test_invalid_params() {
295        let generator = RandomStringGenerator::new();
296        let mut params = HashMap::new();
297        params.insert("length".to_string(), Value::Number(serde_json::Number::from(0)));
298        
299        assert!(generator.validate_params(&params).is_err());
300        
301        params.insert("length".to_string(), Value::Number(serde_json::Number::from(2000)));
302        assert!(generator.validate_params(&params).is_err());
303        
304        params.insert("charset".to_string(), Value::String("".to_string()));
305        assert!(generator.validate_params(&params).is_err());
306    }
307
308    #[test]
309    fn test_duplicate_registration() {
310        let mut registry = CustomGeneratorRegistry::new();
311        let generator1 = RandomStringGenerator::new();
312        let generator2 = RandomStringGenerator::new();
313        
314        assert!(registry.register(generator1).is_ok());
315        assert!(registry.register(generator2).is_err());
316    }
317
318    #[test]
319    fn test_list_generators() {
320        let mut registry = CustomGeneratorRegistry::new();
321        let generator = RandomStringGenerator::new();
322        
323        registry.register(generator).unwrap();
324        
325        let generators = registry.list_generators();
326        assert_eq!(generators.len(), 1);
327        assert_eq!(generators[0].name, "random_string");
328        assert!(!generators[0].description.is_empty());
329        assert!(!generators[0].param_schema.is_empty());
330    }
331}