tweak the annotate wrapper to preserve the function call signature

JBlaschke · Jan 10, 2021 · 8056413 · 8056413
1 parent f36b3f1
commit 8056413
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 4 deletions.
diff --git a/PyNVTX/__init__.py b/PyNVTX/__init__.py
@@ -3,6 +3,7 @@
 
 
 from functools import wraps
+from inspect   import signature
 
 from PyNVTX_backend import *
 from .singleton     import Singleton
@@ -11,7 +12,7 @@
 
 major_version   = 0;
 minor_version   = 3;
-release_version = 0;
+release_version = 1;
 
 NVTX_IDENTIFIER = "NVTX"
 REGISTRY = Registry()
@@ -21,12 +22,13 @@
 def annotate(label):
     def _annotate(func):
         @wraps(func)
-        def wrapped(*args, **kwargs):
+        def wrapper(*args, **kwargs):
             RangePushA(label)
             ret = func(*args, **kwargs)
             RangePop()
             return ret
-        return wrapped
+        wrapper.__signature__ = signature(func)
+        return wrapper
     return _annotate
 
 

diff --git a/setup.py b/setup.py
@@ -180,7 +180,7 @@ def build_extensions(self):
 
     setup(
         name="PyNVTX",
-        version="0.3.0",
+        version="0.3.1",
         author="Johannes Blaschke",
         author_email="[email protected]",
         description="A thin python wrapper for the nvToolsExt (NVTX) library, using pybind11 ... with some bells and whistles thrown in for good measure.",