virtue/generate/
generator.rs

1use super::{GenEnum, GenStruct, GenerateMod, Impl, ImplFor, StreamBuilder, StringOrIdent};
2use crate::parse::{GenericConstraints, Generics};
3use crate::prelude::{Ident, TokenStream};
4
5#[must_use]
6/// The generator is used to generate code.
7///
8/// Often you will want to use [`impl_for`] to generate an `impl <trait_name> for <target_name()>`.
9///
10/// [`impl_for`]: #method.impl_for
11pub struct Generator {
12    name: Ident,
13    generics: Option<Generics>,
14    generic_constraints: Option<GenericConstraints>,
15    stream: StreamBuilder,
16}
17
18impl Generator {
19    pub(crate) fn new(
20        name: Ident,
21        generics: Option<Generics>,
22        generic_constraints: Option<GenericConstraints>,
23    ) -> Self {
24        Self {
25            name,
26            generics,
27            generic_constraints,
28            stream: StreamBuilder::new(),
29        }
30    }
31
32    /// Return the name for the struct or enum that this is going to be implemented on.
33    pub fn target_name(&self) -> Ident {
34        self.name.clone()
35    }
36
37    /// Generate an `impl <target_name>` implementation. See [`Impl`] for more information.
38    ///
39    /// This will default to the type that is associated with this generator. If you need to generate an impl for another type you can use `impl_for_other_type`
40    pub fn r#impl(&mut self) -> Impl<Self> {
41        Impl::with_parent_name(self)
42    }
43
44    /// Generate an `impl <target_name>` implementation. See [`Impl`] for more information.
45    ///
46    /// Alias for [`impl`] which doesn't need a `r#` prefix.
47    ///
48    /// [`impl`]: #method.impl
49    pub fn generate_impl(&mut self) -> Impl<Self> {
50        Impl::with_parent_name(self)
51    }
52
53    /// Generate an `for <trait_name> for <target_name>` implementation. See [ImplFor] for more information.
54    ///
55    /// This will default to the type that is associated with this generator. If you need to generate an impl for another type you can use `impl_trait_for_other_type`
56    pub fn impl_for(&mut self, trait_name: impl Into<String>) -> ImplFor<Self> {
57        ImplFor::new(
58            self,
59            self.name.clone().into(),
60            Some(trait_name.into().into()),
61        )
62    }
63
64    /// Generate an `impl <type_name>` block. See [ImplFor] for more information.
65    /// ```
66    /// # use virtue::prelude::*;
67    /// # let mut generator = Generator::with_name("Baz");
68    /// generator.impl_for_other_type("Foo");
69    ///
70    /// // will output:
71    /// // impl Foo { }
72    /// # generator.assert_eq("impl Foo { }");
73    /// ```
74    pub fn impl_for_other_type(&mut self, type_name: impl Into<StringOrIdent>) -> ImplFor<Self> {
75        ImplFor::new(self, type_name.into(), None)
76    }
77
78    /// Generate an `impl <trait_name> for <type_name>` block. See [ImplFor] for more information.
79    /// ```
80    /// # use virtue::prelude::*;
81    /// # let mut generator = Generator::with_name("Baz");
82    /// generator.impl_trait_for_other_type("Foo", "Bar");
83    ///
84    /// // will output:
85    /// // impl Foo for Bar { }
86    /// # generator.assert_eq("impl Foo for Bar { }");
87    /// ```
88    pub fn impl_trait_for_other_type(
89        &mut self,
90        trait_name: impl Into<StringOrIdent>,
91        type_name: impl Into<StringOrIdent>,
92    ) -> ImplFor<Self> {
93        ImplFor::new(self, type_name.into(), Some(trait_name.into()))
94    }
95
96    /// Generate an `for <..lifetimes> <trait_name> for <target_name>` implementation. See [ImplFor] for more information.
97    ///
98    /// Note:
99    /// - Lifetimes should _not_ have the leading apostrophe.
100    /// - `trait_name` should _not_ have custom lifetimes. These will be added automatically.
101    ///
102    /// ```
103    /// # use virtue::prelude::*;
104    /// # let mut generator = Generator::with_name("Bar");
105    /// generator.impl_for_with_lifetimes("Foo", ["a", "b"]);
106    ///
107    /// // will output:
108    /// // impl<'a, 'b> Foo<'a, 'b> for StructOrEnum { }
109    /// # generator.assert_eq("impl < 'a , 'b > Foo < 'a , 'b > for Bar { }");
110    /// ```
111    ///
112    /// The new lifetimes are not associated with any existing lifetimes. If you want this behavior you can call `.impl_for_with_lifetimes(...).new_lifetimes_depend_on_existing()`
113    ///
114    /// ```
115    /// # use virtue::prelude::*;
116    /// # let mut generator = Generator::with_name("Bar").with_lifetime("a");
117    /// // given a derive on `struct<'a> Bar<'a>`
118    /// generator.impl_for_with_lifetimes("Foo", ["b"]).new_lifetimes_depend_on_existing();
119    ///
120    /// // will output:
121    /// // impl<'a, 'b> Foo<'b> for Bar<'a> where 'b: 'a { }
122    /// # generator.assert_eq("impl < 'b , 'a > Foo < 'b > for Bar < 'a > where 'b : 'a { }");
123    /// ```
124    pub fn impl_for_with_lifetimes<ITER, T>(
125        &mut self,
126        trait_name: T,
127        lifetimes: ITER,
128    ) -> ImplFor<Self>
129    where
130        ITER: IntoIterator,
131        ITER::Item: Into<String>,
132        T: Into<StringOrIdent>,
133    {
134        ImplFor::new(self, self.name.clone().into(), Some(trait_name.into()))
135            .with_lifetimes(lifetimes)
136    }
137
138    /// Generate a struct with the given name. See [`GenStruct`] for more info.
139    pub fn generate_struct(&mut self, name: impl Into<String>) -> GenStruct<Self> {
140        GenStruct::new(self, name)
141    }
142
143    /// Generate an enum with the given name. See [`GenEnum`] for more info.
144    pub fn generate_enum(&mut self, name: impl Into<String>) -> GenEnum<Self> {
145        GenEnum::new(self, name)
146    }
147
148    /// Generate a `mod <name> { ... }`. See [`GenerateMod`] for more info.
149    pub fn generate_mod(&mut self, mod_name: impl Into<String>) -> GenerateMod<Self> {
150        GenerateMod::new(self, mod_name)
151    }
152
153    /// Export the current stream to a file, making it very easy to debug the output of a derive macro.
154    /// This will try to find rust's `target` directory, and write `target/generated/<crate_name>/<name>_<file_postfix>.rs`.
155    ///
156    /// Will return `true` if the file is written, `false` otherwise.
157    ///
158    /// The outputted file is unformatted. Use `cargo fmt -- target/generated/<crate_name>/<file>.rs` to format the file.
159    pub fn export_to_file(&self, crate_name: &str, file_postfix: &str) -> bool {
160        use std::io::Write;
161
162        if let Ok(var) = std::env::var("CARGO_MANIFEST_DIR") {
163            let mut path = std::path::PathBuf::from(var);
164            loop {
165                {
166                    let mut path = path.clone();
167                    path.push("target");
168                    if path.exists() {
169                        path.push("generated");
170                        path.push(crate_name);
171                        if std::fs::create_dir_all(&path).is_err() {
172                            return false;
173                        }
174                        path.push(format!("{}_{}.rs", self.target_name(), file_postfix));
175                        if let Ok(mut file) = std::fs::File::create(path) {
176                            let _ = file.write_all(self.stream.stream.to_string().as_bytes());
177                            return true;
178                        }
179                    }
180                }
181                if let Some(parent) = path.parent() {
182                    path = parent.into();
183                } else {
184                    break;
185                }
186            }
187        }
188        false
189    }
190
191    /// Consume the contents of this generator. This *must* be called, or else the generator will panic on drop.
192    pub fn finish(mut self) -> crate::prelude::Result<TokenStream> {
193        Ok(std::mem::take(&mut self.stream).stream)
194    }
195}
196
197#[cfg(feature = "proc-macro2")]
198impl Generator {
199    /// Create a new generator with the name `name`. This is useful for testing purposes in combination with the `assert_eq` function.
200    pub fn with_name(name: &str) -> Self {
201        Self::new(
202            Ident::new(name, crate::prelude::Span::call_site()),
203            None,
204            None,
205        )
206    }
207    /// Add a lifetime to this generator.
208    pub fn with_lifetime(mut self, lt: &str) -> Self {
209        self.generics
210            .get_or_insert_with(|| Generics(Vec::new()))
211            .push(crate::parse::Generic::Lifetime(crate::parse::Lifetime {
212                ident: crate::prelude::Ident::new(lt, crate::prelude::Span::call_site()),
213                constraint: Vec::new(),
214            }));
215        self
216    }
217    /// Assert that the generated code in this generator matches the given string. This is useful for testing purposes in combination with the `with_name` function.
218    pub fn assert_eq(&self, expected: &str) {
219        assert_eq!(expected, self.stream.stream.to_string());
220    }
221}
222
223impl Drop for Generator {
224    fn drop(&mut self) {
225        if !self.stream.stream.is_empty() && !std::thread::panicking() {
226            eprintln!("WARNING: Generator dropped but the stream is not empty. Please call `.finish()` on the generator");
227        }
228    }
229}
230
231impl super::Parent for Generator {
232    fn append(&mut self, builder: StreamBuilder) {
233        self.stream.append(builder);
234    }
235
236    fn name(&self) -> &Ident {
237        &self.name
238    }
239
240    fn generics(&self) -> Option<&Generics> {
241        self.generics.as_ref()
242    }
243
244    fn generic_constraints(&self) -> Option<&GenericConstraints> {
245        self.generic_constraints.as_ref()
246    }
247}
248
249#[cfg(test)]
250mod test {
251    use proc_macro2::Span;
252
253    use crate::token_stream;
254
255    use super::*;
256
257    #[test]
258    fn impl_for_with_lifetimes() {
259        // No generics
260        let mut generator =
261            Generator::new(Ident::new("StructOrEnum", Span::call_site()), None, None);
262        let _ = generator.impl_for_with_lifetimes("Foo", ["a", "b"]);
263        let output = generator.finish().unwrap();
264        assert_eq!(
265            output
266                .into_iter()
267                .map(|v| v.to_string())
268                .collect::<String>(),
269            token_stream("impl<'a, 'b> Foo<'a, 'b> for StructOrEnum { }")
270                .map(|v| v.to_string())
271                .collect::<String>(),
272        );
273
274        //with simple generics
275        let mut generator = Generator::new(
276            Ident::new("StructOrEnum", Span::call_site()),
277            Generics::try_take(&mut token_stream("<T1, T2>")).unwrap(),
278            None,
279        );
280        let _ = generator.impl_for_with_lifetimes("Foo", ["a", "b"]);
281        let output = generator.finish().unwrap();
282        assert_eq!(
283            output
284                .into_iter()
285                .map(|v| v.to_string())
286                .collect::<String>(),
287            token_stream("impl<'a, 'b, T1, T2> Foo<'a, 'b> for StructOrEnum<T1, T2> { }")
288                .map(|v| v.to_string())
289                .collect::<String>()
290        );
291
292        // with lifetimes
293        let mut generator = Generator::new(
294            Ident::new("StructOrEnum", Span::call_site()),
295            Generics::try_take(&mut token_stream("<'alpha, 'beta>")).unwrap(),
296            None,
297        );
298        let _ = generator.impl_for_with_lifetimes("Foo", ["a", "b"]);
299        let output = generator.finish().unwrap();
300        assert_eq!(
301            output
302                .into_iter()
303                .map(|v| v.to_string())
304                .collect::<String>(),
305            token_stream(
306                "impl<'a, 'b, 'alpha, 'beta> Foo<'a, 'b> for StructOrEnum<'alpha, 'beta> { }"
307            )
308            .map(|v| v.to_string())
309            .collect::<String>()
310        );
311    }
312
313    #[test]
314    fn impl_for_with_trait_generics() {
315        let mut generator = Generator::new(
316            Ident::new("StructOrEnum", Span::call_site()),
317            Generics::try_take(&mut token_stream("<'a>")).unwrap(),
318            None,
319        );
320        let _ = generator.impl_for("Foo").with_trait_generics(["&'a str"]);
321        let output = generator.finish().unwrap();
322        assert_eq!(
323            output
324                .into_iter()
325                .map(|v| v.to_string())
326                .collect::<String>(),
327            token_stream("impl<'a> Foo<&'a str> for StructOrEnum<'a> { }")
328                .map(|v| v.to_string())
329                .collect::<String>(),
330        );
331    }
332
333    #[test]
334    fn impl_for_with_impl_generics() {
335        //with simple generics
336        let mut generator = Generator::new(
337            Ident::new("StructOrEnum", Span::call_site()),
338            Generics::try_take(&mut token_stream("<T1, T2>")).unwrap(),
339            None,
340        );
341        let _ = generator.impl_for("Foo").with_impl_generics(["Bar"]);
342
343        let output = generator.finish().unwrap();
344        assert_eq!(
345            output
346                .into_iter()
347                .map(|v| v.to_string())
348                .collect::<String>(),
349            token_stream("impl<T1, T2, Bar> Foo for StructOrEnum<T1, T2> { }")
350                .map(|v| v.to_string())
351                .collect::<String>()
352        );
353        // with lifetimes
354        let mut generator = Generator::new(
355            Ident::new("StructOrEnum", Span::call_site()),
356            Generics::try_take(&mut token_stream("<'alpha, 'beta>")).unwrap(),
357            None,
358        );
359        let _ = generator.impl_for("Foo").with_impl_generics(["Bar"]);
360        let output = generator.finish().unwrap();
361        assert_eq!(
362            output
363                .into_iter()
364                .map(|v| v.to_string())
365                .collect::<String>(),
366            token_stream("impl<'alpha, 'beta, Bar> Foo for StructOrEnum<'alpha, 'beta> { }")
367                .map(|v| v.to_string())
368                .collect::<String>()
369        );
370    }
371}
virtue/generate/generator.rs

virtue/generate/
generator.rs
