polyhorn_cli/template/
mod.rs

1//! Utilities for quickly building a source tree from a series of templates.
2
3use serde::Serialize;
4use std::fs::{create_dir_all, File};
5use std::io::Write;
6use std::path::Path;
7use tinytemplate::TinyTemplate;
8
9use crate::spec::Spec;
10
11/// Represents the contents of a source file that needs to be generated.
12pub enum SourceFileContents {
13    /// Represents a template. The template will be invoked with the spec
14    /// contained in `Polyhorn.toml`.
15    Template(&'static str),
16
17    /// Copies the contents of a file from the given slice without any string
18    /// interpolation.
19    Copy(&'static [u8]),
20}
21
22/// A single source file that is generated.
23pub struct SourceFile {
24    name: String,
25    contents: SourceFileContents,
26}
27
28/// A source tree that is used to queue files that need to be generated.
29pub struct SourceTree {
30    files: Vec<SourceFile>,
31}
32
33/// Context that is passed to the template engine that is used for generating
34/// source files dynamically.
35#[derive(Serialize)]
36pub struct Context<'a> {
37    spec: &'a Spec,
38}
39
40/// An iterator over the files that a source tree generates one-by-one.
41pub struct SourceTreeIter<'a> {
42    files: &'a [SourceFile],
43    current: usize,
44    context: Context<'a>,
45    destination_path: &'a Path,
46}
47
48/// Represents an error that occurs during source tree generation.
49#[derive(Debug)]
50pub enum Error {
51    /// Contains an error returned by the template engine invocation.
52    Template(tinytemplate::error::Error),
53
54    /// Contains an IO error that is encountered when writing the generated
55    /// source files to disk.
56    IO(std::io::Error),
57}
58
59impl<'a> Iterator for SourceTreeIter<'a> {
60    type Item = Result<(), Error>;
61
62    fn next(&mut self) -> Option<Self::Item> {
63        if self.current == self.files.len() {
64            return None;
65        }
66
67        let file = &self.files[self.current];
68        self.current += 1;
69
70        let mut buf = self.destination_path.to_path_buf();
71        buf.push(file.name.to_owned());
72
73        {
74            let mut buf = buf.clone();
75            buf.pop();
76
77            let _ = create_dir_all(buf);
78        }
79
80        match file.contents {
81            SourceFileContents::Copy(contents) => {
82                let mut file = match File::create(buf) {
83                    Ok(file) => file,
84                    Err(error) => return Some(Err(Error::IO(error))),
85                };
86
87                if let Err(error) = file.write_all(contents) {
88                    return Some(Err(Error::IO(error)));
89                }
90            }
91            SourceFileContents::Template(template) => {
92                let mut engine = TinyTemplate::new();
93
94                if let Err(error) = engine.add_template("template", template) {
95                    return Some(Err(Error::Template(error)));
96                }
97
98                let contents = match engine.render("template", &self.context) {
99                    Ok(contents) => contents,
100                    Err(error) => return Some(Err(Error::Template(error))),
101                };
102
103                let mut file = match File::create(buf) {
104                    Ok(file) => file,
105                    Err(error) => return Some(Err(Error::IO(error))),
106                };
107
108                if let Err(error) = file.write_all(contents.as_bytes()) {
109                    return Some(Err(Error::IO(error)));
110                }
111            }
112        }
113
114        Some(Ok(()))
115    }
116}
117
118impl SourceTree {
119    /// Creates a new empty source tree that can be written to disk.
120    pub fn new() -> SourceTree {
121        SourceTree { files: vec![] }
122    }
123
124    /// Queues a copy of the given contents to a file with the given name.
125    pub fn copy(&mut self, name: &str, contents: &'static [u8]) {
126        self.files.push(SourceFile {
127            name: name.to_owned(),
128            contents: SourceFileContents::Copy(contents),
129        })
130    }
131
132    /// Queues an invocation of the given template that writes the result to a
133    /// file with the given name.
134    pub fn template(&mut self, name: &str, contents: &'static str) {
135        self.files.push(SourceFile {
136            name: name.to_owned(),
137            contents: SourceFileContents::Template(contents),
138        });
139    }
140
141    /// Returns the number of files that are queued for generation.
142    pub fn len(&self) -> usize {
143        self.files.len()
144    }
145
146    /// Returns an iterator that generates the source files of this source tree
147    /// and writes them to the given destination path one-by-one.
148    pub fn generate<'a>(
149        &'a self,
150        spec: &'a Spec,
151        destination_path: &'a Path,
152    ) -> SourceTreeIter<'a> {
153        SourceTreeIter {
154            files: &self.files,
155            current: 0,
156            context: Context { spec },
157            destination_path,
158        }
159    }
160}