subplotlib-derive: Support additional context registration

Signed-off-by: Daniel Silverstone <dsilvers@digital-scurf.org>

1 files changed, 26 insertions, 2 deletions

diff --git a/subplotlib-derive/src/lib.rs b/subplotlib-derive/src/lib.rs
index e2f9a2c..e12ae6a 100644
--- a/subplotlib-derive/src/lib.rs
+++ b/subplotlib-derive/src/lib.rs
@@ -123,7 +123,7 @@ fn check_step_declaration(step: &ItemFn) {
 }
 
 #[throws(Error)]
-fn process_step(input: ItemFn) -> proc_macro2::TokenStream {
+fn process_step(mut input: ItemFn) -> proc_macro2::TokenStream {
     // Processing a step involves constructing a step builder for
     // the function which returns a step object to be passed into the
     // scenario system
@@ -161,6 +161,18 @@ fn process_step(input: ItemFn) -> proc_macro2::TokenStream {
         unreachable!()
     };
 
+    let contexts: Vec<Type> = input
+        .attrs
+        .iter()
+        .filter(|attr| attr.path.is_ident("context"))
+        .map(|attr| {
+            let ty: Type = attr.parse_args()?;
+            Ok(ty)
+        })
+        .collect::<Result<_, Error>>()?;
+
+    input.attrs.retain(|f| !f.path.is_ident("context"));
+
     let fields = input
         .sig
         .inputs
@@ -302,11 +314,23 @@ fn process_step(input: ItemFn) -> proc_macro2::TokenStream {
         }
     };
 
+    let extra_registers: Vec<_> = contexts
+        .iter()
+        .map(|ty| {
+            quote! {
+                scenario.register_context_type::<#ty>();
+            }
+        })
+        .collect();
+
     let register_fn_body = if ty_is_scenariocontext(&contexttype) {
-        quote! {}
+        quote! {
+            #(#extra_registers)*
+        }
     } else {
         quote! {
             scenario.register_context_type::<#contexttype>();
+            #(#extra_registers)*
         }
     };