From cbbeea049540cc3d2b32409f33a1b7bcbe161479 Mon Sep 17 00:00:00 2001
From: Aiden McClelland <me@drbonez.dev>
Date: Fri, 3 May 2024 22:43:48 -0600
Subject: [PATCH] fix custom_display_fn

---
 src/handler/adapters.rs | 11 ++++-------
 tests/handler.rs        |  1 +
 2 files changed, 5 insertions(+), 7 deletions(-)

diff --git a/src/handler/adapters.rs b/src/handler/adapters.rs
index 04f66ec..78fb735 100644
--- a/src/handler/adapters.rs
+++ b/src/handler/adapters.rs
@@ -30,9 +30,9 @@ pub trait HandlerExt<Context: crate::Context>: HandlerFor<Context> + Sized {
     fn with_custom_display_fn<C: crate::Context, F>(
         self,
         display: F,
-    ) -> CustomDisplayFn<F, Self, Context>
+    ) -> CustomDisplayFn<F, Self, C>
     where
-        F: Fn(HandlerArgsFor<Context, Self>, Self::Ok) -> Result<(), Self::Err>;
+        F: Fn(HandlerArgsFor<C, Self>, Self::Ok) -> Result<(), Self::Err>;
     fn with_inherited<Params, InheritedParams, F>(
         self,
         f: F,
@@ -64,12 +64,9 @@ impl<Context: crate::Context, T: HandlerFor<Context> + Sized> HandlerExt<Context
             handler: self,
         }
     }
-    fn with_custom_display_fn<C: crate::Context, F>(
-        self,
-        display: F,
-    ) -> CustomDisplayFn<F, Self, Context>
+    fn with_custom_display_fn<C: crate::Context, F>(self, display: F) -> CustomDisplayFn<F, Self, C>
     where
-        F: Fn(HandlerArgsFor<Context, Self>, Self::Ok) -> Result<(), Self::Err>,
+        F: Fn(HandlerArgsFor<C, Self>, Self::Ok) -> Result<(), Self::Err>,
     {
         CustomDisplayFn {
             _phantom: PhantomData::new(),
diff --git a/tests/handler.rs b/tests/handler.rs
index 1a0d9cf..7128ba2 100644
--- a/tests/handler.rs
+++ b/tests/handler.rs
@@ -129,6 +129,7 @@ fn make_api<C: Context>() -> ParentHandler<C> {
                     ))
                 },
             )
+            .with_custom_display_fn(|_, a| Ok(println!("{a}")))
             .with_call_remote::<CliContext>(),
         )
         .subcommand(