Is there a command to exit a module when imported, like return for a function

Question

When you import a module in python, the module code is "run". Sometimes it is useful to have branching logic in the module such as checking package versions or platform or whatever. Is there a way to exit the entire module execution before hitting the end of the file, something equivalent to early return in a function?

If the module is run as a script you can exit() but that actively raises an exception and kills the whole thing. I just want to say, that's it you are done now, don't run any more code below here.

Basically can I transform this

if not <condition>:
    MY_CONSTANT = 3.14
    class blah():
       ...

    def foo(x):
       ...

    # rest of module....

into

if <condition>:
    return from module

MY_CONSTANT = 3.14
class blah():
   ...

def foo(x):
   ...

# rest of module....

mostly so that I don't have to have lots of code that looks strangely one extra indent level in.

Answer 1

You can create a custom Loader that special-cases e.g. ImportError ⁽¹⁾ as a shortcut to stop module execution. This can be registered via a custom Finder at sys.meta_path.

So if you have the following module to be imported:

# foo.py

x = 1
raise ImportError  # stop module execution here
y = 2

You can use the following finder/loader to import that module. It will be executed until the point where it hits the raise ImportError.

import importlib


class Loader(importlib.machinery.SourceFileLoader):
    def exec_module(self, module):
        try:
            super().exec_module(module)
        except ImportError:  # the module chose to stop executing
            pass


class Finder(importlib.machinery.PathFinder):
    @classmethod
    def find_spec(cls, fullname, path=None, target=None):
        spec = super().find_spec(fullname, path, target)
        if spec is not None:
            spec.loader = Loader(spec.name, spec.origin)  # register the custom loader
        return spec


import sys

sys.meta_path.insert(2, Finder())  # from now on the custom finder will be queried for imports


import foo

print(foo.x)  # prints 1
print(foo.y)  # raises AttributeError

---

(1) Using ImportError to indicate the shortcut obviously has its downsides, such as if your module tries to import something else which doesn't exist, this won't be reported as an error but the module just stops executing. So it's better to use some custom exception instead. I'm just using ImportError for the sake of the example.