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                        if let Ok(mut file) = std::fs::File::create(path) {
179                            let _ = file.write_all(self.stream.stream.to_string().as_bytes());
180                            return true;
181                        }
182                    }
183                }
184                if let Some(parent) = path.parent() {
185                    path = parent.into();
186                } else {
187                    break;
188                }
189            }
190        }
191        false
192    }
193
194    /// Consume the contents of this generator. This *must* be called, or else the generator will panic on drop.
195    pub fn finish(mut self) -> crate::prelude::Result<TokenStream> {
196        Ok(std::mem::take(&mut self.stream).stream)
197    }
198}
199
200#[cfg(feature = "proc-macro2")]
201impl Generator {
202    /// Create a new generator with the name `name`. This is useful for testing purposes in combination with the `assert_eq` function.
203    pub fn with_name(name: &str) -> Self {
204        Self::new(
205            Ident::new(name, crate::prelude::Span::call_site()),
206            None,
207            None,
208        )
209    }
210    /// Add a lifetime to this generator.
211    pub fn with_lifetime(mut self, lt: &str) -> Self {
212        self.generics
213            .get_or_insert_with(|| Generics(Vec::new()))
214            .push(crate::parse::Generic::Lifetime(crate::parse::Lifetime {
215                ident: crate::prelude::Ident::new(lt, crate::prelude::Span::call_site()),
216                constraint: Vec::new(),
217            }));
218        self
219    }
220    /// 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.
221    pub fn assert_eq(&self, expected: &str) {
222        assert_eq!(expected, self.stream.stream.to_string());
223    }
224}
225
226impl Drop for Generator {
227    fn drop(&mut self) {
228        if !self.stream.stream.is_empty() && !std::thread::panicking() {
229            eprintln!("WARNING: Generator dropped but the stream is not empty. Please call `.finish()` on the generator");
230        }
231    }
232}
233
234impl super::Parent for Generator {
235    fn append(&mut self, builder: StreamBuilder) {
236        self.stream.append(builder);
237    }
238
239    fn name(&self) -> &Ident {
240        &self.name
241    }
242
243    fn generics(&self) -> Option<&Generics> {
244        self.generics.as_ref()
245    }
246
247    fn generic_constraints(&self) -> Option<&GenericConstraints> {
248        self.generic_constraints.as_ref()
249    }
250}
251
252#[cfg(test)]
253mod test {
254    use proc_macro2::Span;
255
256    use crate::token_stream;
257
258    use super::*;
259
260    #[test]
261    fn impl_for_with_lifetimes() {
262        // No generics
263        let mut generator =
264            Generator::new(Ident::new("StructOrEnum", Span::call_site()), None, None);
265        let _ = generator.impl_for_with_lifetimes("Foo", ["a", "b"]);
266        let output = generator.finish().unwrap();
267        assert_eq!(
268            output
269                .into_iter()
270                .map(|v| v.to_string())
271                .collect::<String>(),
272            token_stream("impl<'a, 'b> Foo<'a, 'b> for StructOrEnum { }")
273                .map(|v| v.to_string())
274                .collect::<String>(),
275        );
276
277        //with simple generics
278        let mut generator = Generator::new(
279            Ident::new("StructOrEnum", Span::call_site()),
280            Generics::try_take(&mut token_stream("<T1, T2>")).unwrap(),
281            None,
282        );
283        let _ = generator.impl_for_with_lifetimes("Foo", ["a", "b"]);
284        let output = generator.finish().unwrap();
285        assert_eq!(
286            output
287                .into_iter()
288                .map(|v| v.to_string())
289                .collect::<String>(),
290            token_stream("impl<'a, 'b, T1, T2> Foo<'a, 'b> for StructOrEnum<T1, T2> { }")
291                .map(|v| v.to_string())
292                .collect::<String>()
293        );
294
295        // with lifetimes
296        let mut generator = Generator::new(
297            Ident::new("StructOrEnum", Span::call_site()),
298            Generics::try_take(&mut token_stream("<'alpha, 'beta>")).unwrap(),
299            None,
300        );
301        let _ = generator.impl_for_with_lifetimes("Foo", ["a", "b"]);
302        let output = generator.finish().unwrap();
303        assert_eq!(
304            output
305                .into_iter()
306                .map(|v| v.to_string())
307                .collect::<String>(),
308            token_stream(
309                "impl<'a, 'b, 'alpha, 'beta> Foo<'a, 'b> for StructOrEnum<'alpha, 'beta> { }"
310            )
311            .map(|v| v.to_string())
312            .collect::<String>()
313        );
314    }
315
316    #[test]
317    fn impl_for_with_trait_generics() {
318        let mut generator = Generator::new(
319            Ident::new("StructOrEnum", Span::call_site()),
320            Generics::try_take(&mut token_stream("<'a>")).unwrap(),
321            None,
322        );
323        let _ = generator.impl_for("Foo").with_trait_generics(["&'a str"]);
324        let output = generator.finish().unwrap();
325        assert_eq!(
326            output
327                .into_iter()
328                .map(|v| v.to_string())
329                .collect::<String>(),
330            token_stream("impl<'a> Foo<&'a str> for StructOrEnum<'a> { }")
331                .map(|v| v.to_string())
332                .collect::<String>(),
333        );
334    }
335
336    #[test]
337    fn impl_for_with_impl_generics() {
338        //with simple generics
339        let mut generator = Generator::new(
340            Ident::new("StructOrEnum", Span::call_site()),
341            Generics::try_take(&mut token_stream("<T1, T2>")).unwrap(),
342            None,
343        );
344        let _ = generator.impl_for("Foo").with_impl_generics(["Bar"]);
345
346        let output = generator.finish().unwrap();
347        assert_eq!(
348            output
349                .into_iter()
350                .map(|v| v.to_string())
351                .collect::<String>(),
352            token_stream("impl<T1, T2, Bar> Foo for StructOrEnum<T1, T2> { }")
353                .map(|v| v.to_string())
354                .collect::<String>()
355        );
356        // with lifetimes
357        let mut generator = Generator::new(
358            Ident::new("StructOrEnum", Span::call_site()),
359            Generics::try_take(&mut token_stream("<'alpha, 'beta>")).unwrap(),
360            None,
361        );
362        let _ = generator.impl_for("Foo").with_impl_generics(["Bar"]);
363        let output = generator.finish().unwrap();
364        assert_eq!(
365            output
366                .into_iter()
367                .map(|v| v.to_string())
368                .collect::<String>(),
369            token_stream("impl<'alpha, 'beta, Bar> Foo for StructOrEnum<'alpha, 'beta> { }")
370                .map(|v| v.to_string())
371                .collect::<String>()
372        );
373    }
374}
virtue_next/generate/generator.rs

virtue_next/generate/
generator.rs
