Use functools.wraps to preserve handler signature (#3287)

pyro-ppl · Oct 24, 2023 · 23e6470 · 23e6470
1 parent c605f41
commit 23e6470
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/pyro/poutine/block_messenger.py b/pyro/poutine/block_messenger.py
@@ -110,7 +110,7 @@ class BlockMessenger(Messenger):
     :param list hide: list of site names to hide
     :param list expose: list of site names to be exposed while all others hidden
     :param list hide_types: list of site types to be hidden
-    :param lits expose_types: list of site types to be exposed while all others hidden
+    :param list expose_types: list of site types to be exposed while all others hidden
     :returns: stochastic function decorated with a :class:`~pyro.poutine.block_messenger.BlockMessenger`
     """
 

diff --git a/pyro/poutine/handlers.py b/pyro/poutine/handlers.py
@@ -81,6 +81,7 @@
 
 def _make_handler(msngr_cls, module=None):
     def handler_decorator(func):
+        @functools.wraps(func)
         def handler(fn=None, *args, **kwargs):
             if fn is not None and not (
                 callable(fn) or isinstance(fn, collections.abc.Iterable)
@@ -101,7 +102,6 @@ def handler(fn=None, *args, **kwargs):
             )
             + (msngr_cls.__doc__ if msngr_cls.__doc__ else "")
         )
-        handler.__name__ = func.__name__
         if module is not None:
             handler.__module__ = module
         return handler