Extract the code in a "with" statement and remove it from runtime

The context

I'm trying to create an "environment" context manager. Think of it as choosing to execute some code locally or remotely depending on a parameter of the context manager:

with PythonExecutor(env="local"):
    x = 1
    assert x == 1

would run that code in-process. However, changing the env parameter to "remote" would connect to SSH and execute the code remotely.

Thanks to this StackOverflow question, I managed to extract the code within the with block as a string in the __exit__ method and the SSH part is trivial (and irrelevant for that question).

The question

How can I prevent the code within the with block to run in-process? Context managers always follow:

1. Calling __enter__
2. Executing the code within the with block
3. Calling __exit__

This means that even if I choose "remote" execution, the code will be executed remotely in __enter__ or __exit__, but it will still be executed locally. In other words, is there some way to skip step 2? I started looking into runtime bytecode manipulation but it's getting a bit hairy…

Other solutions to the original issue (running code in different environments in an elegant way) are welcome too 🙂

It's a bit hacky and requires changing the code of the with block slightly but you could make your __enter__ method return a function that raises an error when env == 'remote'. Then on the remote case you'll get a local error and then handle everything else in the __exit__ block.

class PythonExecutor:

    def __init__(self, env):
        self.env = env

    def __enter__(self):
        def switch():
            if self.env == 'remote':
                raise Exception # Probably some custom exception for only this purpose
        return switch

    def __exit__(self, exctype, excinst, exctb):
        # do ssh stuff here depending on exctype
        ...


with PythonExecutor(env='remote') as switch:
    switch()
    print('hello')