From e101cc4f44e5b126e3826cd2959b64f6d624e26f Mon Sep 17 00:00:00 2001
From: Alex Willmer <alex@moreati.org.uk>
Date: Sat, 23 Apr 2022 18:24:41 +0100
Subject: [PATCH] mitogen.utils: Preserve docstring of functions decorated
 @with_router

Co-authored-by: Rezart Qelibari <gast-kontakt+mitogen@astzweig.de>

Replaces #837
Fixes #836
---
 docs/changelog.rst  | 1 +
 mitogen/utils.py    | 6 ++----
 tests/utils_test.py | 5 +++++
 3 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/docs/changelog.rst b/docs/changelog.rst
index 8feb985b..513ce8e6 100644
--- a/docs/changelog.rst
+++ b/docs/changelog.rst
@@ -24,6 +24,7 @@ v0.3.3.dev0
 * :gh:issue:`906` Support packages dynamically inserted into sys.modules, e.g. `distro` >= 1.7.0 as `ansible.module_utils.distro`.
 * :gh:issue:`918` Support Python 3.10
 * :gh:issue:`920` Support Ansible :ans:conn:`~podman` connection plugin
+* :gh:issue:`836` :func:`mitogen.utils.with_router` decorator preserves the docstring in addition to the name.
 
 
 v0.3.2 (2022-01-12)
diff --git a/mitogen/utils.py b/mitogen/utils.py
index 4e7a59a9..71f7c35f 100644
--- a/mitogen/utils.py
+++ b/mitogen/utils.py
@@ -29,6 +29,7 @@
 # !mitogen: minify_safe
 
 import datetime
+import functools
 import logging
 import os
 import sys
@@ -171,12 +172,9 @@ def with_router(func):
 
         do_stuff(blah, 123)
     """
+    @functools.wraps(func)
     def wrapper(*args, **kwargs):
         return run_with_router(func, *args, **kwargs)
-    if mitogen.core.PY3:
-        wrapper.func_name = func.__name__
-    else:
-        wrapper.func_name = func.func_name
     return wrapper
 
 
diff --git a/tests/utils_test.py b/tests/utils_test.py
index 8ff730ef..6c71d33c 100644
--- a/tests/utils_test.py
+++ b/tests/utils_test.py
@@ -12,6 +12,7 @@ def func0(router):
 
 @mitogen.utils.with_router
 def func(router):
+    "Docstring of func"
     return router
 
 
@@ -31,6 +32,10 @@ class WithRouterTest(testlib.TestCase):
         self.assertIsInstance(router, mitogen.master.Router)
         self.assertFalse(testlib.threading__thread_is_alive(router.broker._thread))
 
+    def test_with_broker_preserves_attributes(self):
+        self.assertEqual(func.__doc__, 'Docstring of func')
+        self.assertEqual(func.__name__, 'func')
+
 
 class Dict(dict): pass
 class List(list): pass