Added functools.wrap for decorators (#55)

gaogaotiantian · Mar 15, 2022 · dffa62c · dffa62c
1 parent b4808f0
commit dffa62c
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 2 deletions.
diff --git a/src/objprint/decorator.py b/src/objprint/decorator.py
@@ -2,6 +2,9 @@
 # For details: https://github.com/gaogaotiantian/objprint/blob/master/NOTICE.txt
 
 
+import functools
+
+
 def add_objprint(orig_class=None, format="string", **kwargs):
 
     from . import _objprint
@@ -19,9 +22,9 @@ def __str__(self):
 
     if orig_class is None:
         def wrapper(cls):
-            cls.__str__ = __str__
+            cls.__str__ = functools.wraps(cls.__str__)(__str__)
             return cls
         return wrapper
     else:
-        orig_class.__str__ = __str__
+        orig_class.__str__ = functools.wraps(orig_class.__str__)(__str__)
         return orig_class
diff --git a/tests/test_decorator.py b/tests/test_decorator.py
@@ -88,6 +88,12 @@ def test_wrapper(self):
 
         self.assertEqual(output, expected)
 
+    def test_signature(self):
+        obj = DecoratedClass()
+        self.assertIn("DecoratedClass", str(obj.__str__))
+        obj = TestWrapper()
+        self.assertIn("TestWrapper", str(obj.__str__))
+
     def test_json(self):
         with io.StringIO() as buf, redirect_stdout(buf):
             print(DecoratedJsonCls())