prmt/
template.rs

1use crate::error::Result;
2use crate::executor::validate_module_occurrences;
3use crate::module_trait::ModuleContext;
4use crate::parser::{Token, parse};
5use crate::registry::ModuleRegistry;
6use crate::style::{AnsiStyle, ModuleStyle, global_no_color};
7use is_terminal::IsTerminal;
8
9/// A parsed template intended for a single render.
10pub struct Template<'a> {
11    tokens: Vec<Token<'a>>,
12    estimated_size: usize,
13}
14
15impl<'a> Template<'a> {
16    /// Parse a template string.
17    #[inline]
18    pub fn new(template: &'a str) -> Self {
19        let tokens = parse(template);
20        let estimated_size = template.len() + (template.len() / 2) + 128;
21        Self {
22            tokens,
23            estimated_size,
24        }
25    }
26
27    /// Render the template with the given registry and context
28    pub fn render(&self, registry: &ModuleRegistry, context: &ModuleContext) -> Result<String> {
29        validate_module_occurrences(&self.tokens)?;
30
31        let mut output = String::with_capacity(self.estimated_size);
32
33        let no_color = global_no_color() || !IsTerminal::is_terminal(&std::io::stdout());
34
35        for token in &self.tokens {
36            match token {
37                Token::Text(text) => {
38                    output.push_str(text);
39                }
40                Token::Placeholder(params) => {
41                    let module = registry.get(&params.module).ok_or_else(|| {
42                        crate::error::PromptError::UnknownModule(params.module.to_string())
43                    })?;
44
45                    if let Some(text) = module.render(&params.format, context)?
46                        && !text.is_empty()
47                    {
48                        let has_prefix = !params.prefix.is_empty();
49                        let has_suffix = !params.suffix.is_empty();
50                        let styled = !params.style.is_empty() && !no_color;
51
52                        if styled {
53                            let style = AnsiStyle::parse(&params.style).map_err(|error| {
54                                crate::error::PromptError::StyleError {
55                                    module: params.module.to_string(),
56                                    error,
57                                }
58                            })?;
59
60                            style.write_start_codes(&mut output, context.shell);
61                            if has_prefix {
62                                output.push_str(&params.prefix);
63                            }
64                            output.push_str(&text);
65                            if has_suffix {
66                                output.push_str(&params.suffix);
67                            }
68                            style.write_reset(&mut output, context.shell);
69                        } else {
70                            if has_prefix {
71                                output.push_str(&params.prefix);
72                            }
73                            output.push_str(&text);
74                            if has_suffix {
75                                output.push_str(&params.suffix);
76                            }
77                        }
78                    }
79                }
80            }
81        }
82
83        Ok(output)
84    }
85
86    /// Get an iterator over the tokens in this template
87    pub fn tokens(&self) -> impl Iterator<Item = &Token<'a>> {
88        self.tokens.iter()
89    }
90
91    /// Get the number of tokens in this template
92    pub fn token_count(&self) -> usize {
93        self.tokens.len()
94    }
95}
96
97#[cfg(test)]
98mod tests {
99    use super::*;
100    use crate::error::PromptError;
101    use crate::module_trait::Module;
102    use std::sync::Arc;
103    use std::sync::atomic::{AtomicUsize, Ordering};
104
105    struct CountingModule {
106        calls: Arc<AtomicUsize>,
107    }
108
109    impl Module for CountingModule {
110        fn render(&self, _format: &str, _context: &ModuleContext) -> Result<Option<String>> {
111            self.calls.fetch_add(1, Ordering::SeqCst);
112            Ok(Some("value".to_string()))
113        }
114    }
115
116    #[test]
117    fn duplicate_singleton_is_rejected_before_rendering() {
118        let calls = Arc::new(AtomicUsize::new(0));
119        let mut registry = ModuleRegistry::new();
120        registry.register(
121            "test",
122            Arc::new(CountingModule {
123                calls: Arc::clone(&calls),
124            }),
125        );
126        let template = Template::new("{test}{test}");
127
128        let error = template
129            .render(&registry, &ModuleContext::default())
130            .expect_err("duplicate module should fail");
131
132        assert!(matches!(
133            error,
134            PromptError::DuplicateModule(module) if module == "test"
135        ));
136        assert_eq!(calls.load(Ordering::SeqCst), 0);
137    }
138}
prmt/template.rs

prmt/
template.rs
