from nilmdb.utils.printf import *
import sys
import inspect
import decorator

def must_close(errorfile = sys.stderr, wrap_verify = False):
    """Class decorator that warns on 'errorfile' at deletion time if
    the class's close() member wasn't called.

    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 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__)
            except AttributeError:
                orig = lambda x: None
            if is_method_or_function(orig):
                setattr(cls, wrapper.__name__,
                        decorator.decorator(wrapper, orig))

        @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

        @wrap_class_method
        def __del__(orig, self, *args, **kwargs):
            try:
                if "_must_close" in self.__dict__:
                    fprintf(errorfile, "error: %s.close() wasn't called!\n",
                            self.__class__.__name__)
                return orig(self, *args, **kwargs)
            except:
                pass

        @wrap_class_method
        def close(orig, self, *args, **kwargs):
            if "_must_close" in self.__dict__:
                del self._must_close
            return orig(self, *args, **kwargs)

        # Optionally wrap all other functions
        def verifier(orig, self, *args, **kwargs):
            if ("_must_close" not in self.__dict__ and
                "_must_close_initialized" in self.__dict__):
                raise AssertionError("called " + str(orig) + " after close")
            return orig(self, *args, **kwargs)
        if wrap_verify:
            for (name, method) in inspect.getmembers(cls, is_method_or_function):
                # Skip some methods
                if name in [ "__del__", "__init__" ]:
                    continue
                # Set up wrapper
                if inspect.ismethod(method):
                    func = method.__func__
                else:
                    func = method
                setattr(cls, name, decorator.decorator(verifier, func))

        return cls
    return class_decorator
