A context to exit deeply nested loops

A context to exit deeply nested loops - Printable Version

+- Python Forum (https://python-forum.io)
+-- Forum: General (https://python-forum.io/forum-1.html)
+--- Forum: Code sharing (https://python-forum.io/forum-5.html)
+--- Thread: A context to exit deeply nested loops (/thread-18243.html)

A context to exit deeply nested loops - Gribouillis - May-10-2019

The break statement doesn't allow to exit several nested loops. It is a simple exercise to develop a block() context to do so. Here is an attempt

from contextlib import contextmanager

@contextmanager
def block():
    """Block of code context allowing to exit deeply nested loops
    
    Usage:
    
        with block() as b:
            ...
                    # somewhere in a nested loop
                    b.exit()
            ...
            
    """
    tok = BlockToken()
    try:
        yield tok
    except EndOfBlock as exc:
        if exc.args[0] is not tok:
            raise

class EndOfBlock(RuntimeError):
    """Helper exception type for the block() context"""

class BlockToken:
    """Helper class for block() context"""
    def exit(self):
        raise EndOfBlock(self)

def func(x):
    """Test function for the block() context"""
    result = ''
    with block() as outer:
        result += 'A'
        for i in range(4):
            result += 'B'
            if x > 0 and i > 2:
                outer.exit()
            result += 'D'
            with block() as inner:
                result += 'C'
                if x == 1:
                    outer.exit()
                elif x == 2:
                    inner.exit()
                result += 'E'
            result += 'F'
        result += 'G'
    result += 'H'
    return result

if __name__ == '__main__':
    assert func(0) == "ABDCEFBDCEFBDCEFBDCEFGH"
    assert func(1) == "ABDCH"
    assert func(2) == "ABDCFBDCFBDCFBH"
    print('SUCCESS')

An interesting feature here is that the BlockToken instance can be used as argument for functions called within the block. For example you could write

def do_something(session):
    if input('Do you want to quit?').strip().lower() in ('y', 'yes'):
        session.exit()
    for i in range(5):
        print(i)

with block() as user_session:
    do_something(user_session)
    print("If we're here, the user chose to stay!")

RE: A context to exit deeply nested loops - Yoriz - May-10-2019

If the nested loop is in a function, return can be used instead of break to exit.

RE: A context to exit deeply nested loops - Gribouillis - May-10-2019

Yoriz Wrote:if the nested loop is in a function, return can be used instead of break to exit.

Yes, but in this case, session.exit() exits more than the function. It exits from a block of code that contained the call to the function.

RE: A context to exit deeply nested loops - Gribouillis - May-19-2019

In this second version, we avoid using the @contextmanager decorator. Instead, block is now a class. The advantage of this is that one can subclass block to create exitable contexts with more data, such as a user session with a user name, a user device, etc

__version__ = '2109.05.19'

class block:
    """Block of code context allowing to exit deeply nested loops
    
    Usage:
    
        with block() as b:
            ...
                    # somewhere in a nested loop
                    b.exit()
            ...
            
    This type can be subclassed to create exitable contexts with
    useful data, for example
    
        class UserSession(block)
            def __init__(self, user_name):
                self.user_name = user_name
                
        with UserSession('Doe') as session:
            ...
                        session.exit()
            ...
            
    """
    
    def __enter__(self):
        return self

    def exit(self):
        raise EndOfBlock(self)
    
    def __exit__(self, exc_type, exc_value, ex_tb):
        if exc_type is EndOfBlock and exc_value.args[0] is self:
            # exception isn't propagated if __exit__ returns true value
            return True

class EndOfBlock(RuntimeError):
    """Helper exception type for the block() context"""


if __name__ == '__main__':
    # Example code for blocks

    class UserSession(block):
        def __init__(self, user_name):
            super().__init__()
            self.user_name = user_name


        def __enter__(self):
            print("User {}'s session starting now!".format(self.user_name))
            return super().__enter__()

        def __exit__(self, *args):
            print('User {} is leaving.'.format(self.user_name))
            return super().__exit__(*args)

    def func(x):
        """Test function for the block() context"""
        result = ''
        with UserSession('Doe') as outer:
            result += 'A'
            for i in range(4):
                result += 'B'
                if x > 0 and i > 2:
                    outer.exit()
                result += 'D'
                with block() as inner:
                    result += 'C'
                    if x == 1:
                        outer.exit()
                    elif x == 2:
                        inner.exit()
                    result += 'E'
                result += 'F'
            result += 'G'
        result += 'H'
        return result

    assert func(0) == "ABDCEFBDCEFBDCEFBDCEFGH"
    assert func(1) == "ABDCH"
    assert func(2) == "ABDCFBDCFBDCFBH"
    print('SUCCESS')

Output:User Doe's session starting now!
User Doe is leaving.
User Doe's session starting now!
User Doe is leaving.
User Doe's session starting now!
User Doe is leaving.
SUCCESS