Go-style 'defer' in Python

defer.py

import functools
import socket
import sys
import threading
import types
import unittest

class OnException(object):
    '''Wrapper for cleanup actions only to be executed upon exception'''

    __slots__ = '_fn',

    def __init__(self, fn):
        self._fn = fn

    def __call__(self):
        return self._fn()

    def __repr__(self):
        return '<defer.OnException(%r) object at %x>' % (self._fn, id(self))

on_exception = OnException


class _Return(BaseException):
    '''Hack to allow to return values from inside a generator'''

    __slots__ = '_value',

    def __init__(self, value):
        self._value = value


def return_(value):
    raise _Return(value)


def with_defer(fn):
    @functools.wraps(fn)
    def wrapped(*args, **kwargs):
        cleanups = []
        result = None
        excepted = False
        exc_info = None

        try:
            gen = fn(*args, **kwargs)
            if not isinstance(gen, types.GeneratorType):
                return gen

            for cleanup in gen:
                cleanups.append(cleanup)
        except _Return as exc:
            result = exc._value
        except:
            exc_info = sys.exc_info()
            excepted = True

        for cleanup in cleanups:
            try:
                if isinstance(cleanup, OnException):
                    if excepted:
                        print 'Cleanup, running:', cleanup
                        cleanup()
                    else:
                        print 'No exception, skipping', cleanup
                else:
                    print 'Cleanup, running:', cleanup
                    cleanup()
            except:
                logging.exception('Error during cleanup')

        if excepted:
            raise exc_info[0], exc_info[1], exc_info[2]
        else:
            return result

    return wrapped


# Demo
count = 0

class TestFoo(unittest.TestCase):
    @with_defer
    def setUp(self):
        global count

        print 'setUp'

        self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        # Only close when this method errors out, otherwise cleaned up by
        # tearDown
        yield on_exception(self.socket.close)

        self.lock = threading.Lock()

        self.lock.acquire()
        # Always release the lock when this method exits
        yield self.lock.release

        if count > 0:
            1 / 0

        count += 1

        print 'setUp done'

        # I know setUp isn't supposed to return something, but this shows it's
        # possible
        return_(1)

    def tearDown(self):
        print 'tearDown'
        self.socket.close()

    def test_foo(self):
        pass

    def test_bar(self):
        pass

gistfile1.txt

$ python -m unittest -v defer.TestFoo
test_bar (defer.TestFoo) ... setUp
setUp done
No exception, skipping <defer.OnException(<bound method _socketobject.close of <socket._socketobject object at 0x7fee12df6de0>>) object at 7fee12dd4248>
Cleanup, running: <built-in method release of thread.lock object at 0x7fee1b775210>
tearDown
ok
test_foo (defer.TestFoo) ... setUp
Cleanup, running: <defer.OnException(<bound method _socketobject.close of <socket._socketobject object at 0x7fee12df6e50>>) object at 7fee12dd4088>
Cleanup, running: <built-in method release of thread.lock object at 0x7fee1b7751f0>
ERROR

======================================================================
ERROR: test_foo (defer.TestFoo)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "defer.py", line 51, in wrapped
    for cleanup in gen:
  File "defer.py", line 103, in setUp
    1 / 0
ZeroDivisionError: integer division or modulo by zero

----------------------------------------------------------------------
Ran 2 tests in 0.001s

FAILED (errors=1)

-> https://stackoverflow.com/a/53069630/9456786