Update mustclose and test_mustclose for Python 3 (this was hard!)

nilmdb/utils/mustclose.py

@@ -12,12 +12,17 @@ def must_close(errorfile = sys.stderr, wrap_verify = False):
     already been called."""
     def class_decorator(cls):
 
+        def is_method_or_function(x):
+            return inspect.ismethod(x) or inspect.isfunction(x)
+
         def wrap_class_method(wrapper):
             try:
-                orig = getattr(cls, wrapper.__name__).__func__
+                orig = getattr(cls, wrapper.__name__)
-            except Exception:
+            except AttributeError:
                 orig = lambda x: None
-            setattr(cls, wrapper.__name__, decorator.decorator(wrapper, orig))
+            if is_method_or_function(orig):
+                setattr(cls, wrapper.__name__,
+                        decorator.decorator(wrapper, orig))
 
         @wrap_class_method
         def __init__(orig, self, *args, **kwargs):
@@ -49,16 +54,16 @@ def must_close(errorfile = sys.stderr, wrap_verify = False):
                 raise AssertionError("called " + str(orig) + " after close")
             return orig(self, *args, **kwargs)
         if wrap_verify:
-            for (name, method) in inspect.getmembers(cls, inspect.ismethod):
+            for (name, method) in inspect.getmembers(cls, is_method_or_function):
-                # Skip class methods
-                if method.__self__ is not None:
-                    continue
                 # Skip some methods
                 if name in [ "__del__", "__init__" ]:
                     continue
                 # Set up wrapper
-                setattr(cls, name, decorator.decorator(verifier,
+                if inspect.ismethod(method):
-                                                       method.__func__))
+                    func = method.__func__
+                else:
+                    func = method
+                setattr(cls, name, decorator.decorator(verifier, func))
 
         return cls
     return class_decorator

tests/test_mustclose.py

@@ -71,6 +71,7 @@ class TestMustClose(object):
 
         # No error
         err.truncate(0)
+        err.seek(0)
         y = Foo("bye")
         y.close()
         del y
@@ -82,6 +83,7 @@ class TestMustClose(object):
 
         # Verify function calls when wrap_verify is True
         err.truncate(0)
+        err.seek(0)
         z = Bar()
         eq_(inspect.getargspec(z.blah),
             inspect.ArgSpec(args = ['self', 'arg'],
@@ -90,7 +92,7 @@ class TestMustClose(object):
         z.close()
         with assert_raises(AssertionError) as e:
             z.blah("hello")
-        in_("called <function blah at 0x", str(e.exception))
+        in_("called <function Bar.blah at 0x", str(e.exception))
         in_("> after close", str(e.exception))
         # Since the most recent assertion references 'z',
         # we need to raise another assertion here so that
@@ -107,6 +109,7 @@ class TestMustClose(object):
 
         # Class with missing methods
         err.truncate(0)
+        err.seek(0)
         w = Baz()
         w.close()
         del w