Monkey patch context manager

MonkeyPatch.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-


class MonkeyPatch(object):
    def __init__(self, original, patch):
        self.original = original
        self.patch = patch
        self.cache = None

        import sys
        self.module_name = self.original.__module__
        self.attr_name = self.original.__name__
        self.module = sys.modules[self.module_name]

    def __enter__(self):
        self.cache = self.original
        setattr(self.module, self.attr_name, self.patch)
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        setattr(self.module, self.attr_name, self.cache)


class ValuePatch(object):
    """ Patch a variable with given value.
        oringnal must be a string.
        patch can be any thing.

        Usage:
            # assume Module.var = "origin"
            with ValuePatch("Module.var", "patch"):
                print Module.var  # patch
            print Module.var  # origin
    """
    def __init__(self, original, patch):
        self.original = original
        self.patch = patch
        self.cache = None

        lst = self.original.split(".")
        self.module_name = ".".join(lst[:-1])
        self.attr_name = "".join(lst[-1])
        self.module = eval(self.module_name)

    def __enter__(self):
        self.cache = eval(self.original)
        setattr(self.module, self.attr_name, self.patch)
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        setattr(self.module, self.attr_name, self.cache)