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

virtue_next/generate/
generator.rs
