Cleanup nilmdb.utils.must_close a bit

nilmdb/utils/mustclose.py

@@ -12,15 +12,12 @@ def must_close(errorfile = sys.stderr, wrap_verify = False):
     already been called."""
     def class_decorator(cls):
 
-        # 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 wrap_class_method(wrapper):
+            try:
+                orig = getattr(cls, wrapper.__name__).im_func
+            except:
+                orig = lambda x: None
+            setattr(cls, wrapper.__name__, decorator.decorator(wrapper, orig))
 
         @wrap_class_method
         def __init__(orig, self, *args, **kwargs):