virtue_next/generate/
generator.rs

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

virtue_next/generate/
generator.rs
