Context manager version of setattr that will do a save/restore (i.e., save previous value, restore it on context exit)

saverestore.py

class SetattrThenRestore:
    """Save the old value of an attribute, set it, then restore it.

    with SetattrThenRestore(obj, attrname, val):
        ...bunch of code here...

    is SOMEWHAT equivalent to:

        oldattrval = getattr(obj, attrname)
        setattr(obj, attrname, val)
        ...bunch of code here...
        setattr(obj, attrname, oldattrval)

    but because this is a context manager the oldattrval gets restored no
    matter what happens in the "bunch of code" (including, of course,
    exceptions and return statements).

    Optional keyword-only argument 'initval' can be specified to force the
    attribute to have a init value IF IT DID NOT ALREADY EXIST in obj.
    Thus, for example:

        class C:
            pass

        foo = C()
        with SetattrThenRestore(foo, 'clown', 'bozo', initval='sad'):
            print(foo.clown)
        print(foo.clown)

    will print:
        bozo
        sad

    whereas "with SetattrThenRestore(foo, 'clown', 'bozo'): ... " will cause
    an AttributeError exception because foo has no initialized clown attr.

    If optional keyword-only argument 'delete' is True (default: False) then
    if the attribute did NOT exist, it will be "restored" on context exit
    by deleting it (restoring the object to the "no attribute" condition).

    It is an error to specify both initval and delete=True.
    """

    NOTGIVEN = object()

    def __init__(self, obj, attrname, tmpval, *,
                 initval=NOTGIVEN, delete=False):
        self.obj = obj
        self.attrname = attrname
        self.tmpval = tmpval
        self.delete = delete
        if initval is not self.NOTGIVEN:
            if delete:
                raise ValueError("cannot specify initval and delete")
            self.initval = initval

    def __enter__(self):
        try:
            self.oldvalue = getattr(self.obj, self.attrname)
        except AttributeError:
            if hasattr(self, 'initval'):
                self.oldvalue = self.initval
            elif not self.delete:
                raise

        setattr(self.obj, self.attrname, self.tmpval)

    def __exit__(self, exc_type, exc_val, exc_tb):
        try:
            setattr(self.obj, self.attrname, self.oldvalue)
        except AttributeError:
            delattr(self.obj, self.attrname)


if __name__ == "__main__":
    import unittest

    class C:
        pass

    class TestMethods(unittest.TestCase):
        def test_missingattribute(self):
            with self.assertRaises(AttributeError):
                foo = C()      # has no attributes
                with SetattrThenRestore(foo, 'a', 17):
                    pass

        def test_init_and_delete(self):
            with self.assertRaises(ValueError):
                foo = C()
                # can't set both initval and delete
                with SetattrThenRestore(foo, 'a', 17, initval=1, delete=True):
                    pass

        def test_initval(self):
            foo = C()
            initval = 6
            testval = 17
            with SetattrThenRestore(foo, 'a', testval, initval=initval):
                self.assertEqual(foo.a, testval)
            self.assertEqual(foo.a, initval)

        def test_deleteattr(self):
            foo = C()
            testval = 17
            with SetattrThenRestore(foo, 'a', testval, delete=True):
                self.assertEqual(foo.a, testval)

            with self.assertRaises(AttributeError):
                _ = foo.a

        def test_saverestore(self):
            foo = C()
            prev = 17
            tmpvalue = 42

            foo.a = prev
            with SetattrThenRestore(foo, 'a', tmpvalue):
                self.assertEqual(foo.a, tmpvalue)
            self.assertEqual(foo.a, prev)

    unittest.main()