MemoryLimit Context Manager / example

import logging
import resource
import sys


class MemoryLimit:

    def __init__(self, soft, hard, log):
        self.soft = soft
        self.hard = hard

    def __enter__(self):
        self.set_limit()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        if isinstance(exc_val, MemoryError):
            log.critical('MemoryError')
            return True

    def set_limit(self):
        resource.setrlimit(resource.RLIMIT_AS, (self.soft, self.hard))


def main():
    list(range(int(1e9)))
    # 1/0
    return 0


if __name__ == '__main__':
    logging.basicConfig()
    log = logging.getLogger()
    MiB = 1 * 1024 ** 2
    with MemoryLimit(5 * MiB, 6 * MiB, log):
        sys.exit(main())
    print('Cleanup Code')
    sys.exit(1)