Fix issue where mustclose decorator doesn't maintain argspec

nilmdb/utils/mustclose.py

@@ -1,6 +1,7 @@
 from nilmdb.utils.printf import *
 import sys
 import inspect
+import decorator
 import functools
 
 def must_close(errorfile = sys.stderr, wrap_verify = False):
@@ -10,39 +11,36 @@ def must_close(errorfile = sys.stderr, wrap_verify = False):
     If 'wrap_verify' is True, every class method is wrapped with a
     verifier that will raise AssertionError if the .close() method has
     already been called."""
-    def decorator(cls):
-        def dummy(*args, **kwargs):
-            pass
-        if "__init__" not in cls.__dict__:
-            cls.__init__ = dummy
-        if "__del__" not in cls.__dict__:
-            cls.__del__ = dummy
-        if "close" not in cls.__dict__:
-            cls.close = dummy
+    def class_decorator(cls):
 
-        orig_init = cls.__init__
-        orig_del = cls.__del__
-        orig_close = cls.close
+        # Helper to replace a class method with a wrapper function,
+        # while maintaining argument specs etc.
+        def wrap_class_method(wrapper_func):
+            method = wrapper_func.__name__
+            if method in cls.__dict__:
+                orig = getattr(cls, method).im_func
+            else:
+                orig = lambda self: None
+            setattr(cls, method, decorator.decorator(wrapper_func, orig))
 
-        def __init__(self, *args, **kwargs):
-            ret = orig_init(self, *args, **kwargs)
+        @wrap_class_method
+        def __init__(orig, self, *args, **kwargs):
+            ret = orig(self, *args, **kwargs)
             self.__dict__["_must_close"] = True
             self.__dict__["_must_close_initialized"] = True
             return ret
 
-        def __del__(self):
+        @wrap_class_method
+        def __del__(orig, self, *args, **kwargs):
             if "_must_close" in self.__dict__:
                 fprintf(errorfile, "error: %s.close() wasn't called!\n",
                         self.__class__.__name__)
-            return orig_del(self)
+            return orig(self, *args, **kwargs)
 
-        def close(self, *args, **kwargs):
+        @wrap_class_method
+        def close(orig, self, *args, **kwargs):
             del self._must_close
-            return orig_close(self)
-
-        cls.__init__ = __init__
-        cls.__del__ = __del__
-        cls.close = close
+            return orig(self, *args, **kwargs)
 
         # Optionally wrap all other functions
         def make_verifier(func):
@@ -65,4 +63,4 @@ def must_close(errorfile = sys.stderr, wrap_verify = False):
                 setattr(cls, name, make_verifier(method))
 
         return cls
-    return decorator
+    return class_decorator