Python Testing a CLI Script Without .py

I often write little CLI programs in Python, where I just want to call it “script” rather than “script.py”. I also want to be able to test functions in these scripts. It gets tricky if you want to test code from a file without a “.py” suffix.

Immediately below is a “drag and drop” fully implemented and documented example, with some further discussion down below that.

Given the following script called “square”:

#!/usr/bin/env python3

def square(x: int) ->  int:
    return x*x

if __name__ == '__main__':
    import sys
    i = int(sys.argv[1])
    print(f'{i} squared is {square(i)}')

Create a “test_square.py” file with the following:

#!/usr/bin/env python3

from typing import List
from types import ModuleType


def import_script_as_module(module_name: str, paths_to_try: List[str]) -> ModuleType:
    """
    Imports a Python script as a module, whether it ends in ".py" or not.
    Given the name of a module to import, and a list of absolute or relative path names
    (including the filename), import the module.  The module is set up so that it can
    later be imported, but a reference to the module is also returned.

    Args:
        module_name (str): The name of the module to import.
                This doesn't have to match the filename.
        paths_to_try (List[str]): A list of file paths to look for the file to load.
                This can be absolute or relative paths to the file, the first file that
                exists is used.

    Returns:
        Module: A reference to the imported module.

    Raises:
        FileNotFoundError: If the module file is not found in any of the specified directory paths.
        ImportError: If there are issues importing the module, such as invalid Python syntax in the module file.

    Example:
        my_module = import_script_as_module("my_module", ["my_module", "../my_module"])

        # Now you can either directly use "my_module"
        my_module.function()

        # Or you can later import it:
        import my_module
    """
    from pathlib import Path
    import os

    for try_filename in paths_to_try:
        if os.path.exists(try_filename):
            module_filename = Path(try_filename).resolve()
            break
    else:
        raise FileNotFoundError(f"Unable to find '{module_name}' module to import")

    from importlib.util import spec_from_loader, module_from_spec
    from importlib.machinery import SourceFileLoader
    import sys

    spec = spec_from_loader(
        module_name, SourceFileLoader(module_name, str(module_filename))
    )
    if spec is None:
        raise ImportError("Unable to spec_from_loader() the module, no error returned.")
    module = module_from_spec(spec)
    spec.loader.exec_module(module)
    sys.modules["up"] = module

    return module


square = import_script_as_module("square", ["square", "../square"])

import unittest


class TestSquare(unittest.TestCase):
    def test_simple(self):
        self.assertEqual(4, square.square(2))

The core of this code to import the “script” as a module is:

#!/usr/bin/env python3

from importlib.util import spec_from_loader, module_from_spec

▎ from importlib.machinery import SourceFileLoader ▎ import sys

script_full_path = Path("script").resolve()
script_str = 'script'

▎ spec = spec_from_loader(script_str, SourceFileLoader(script_str, str(script_full_path)))  script = module_from_spec(spec)  spec.loader.exec_module(script) ▎ sys.modules[script_str] = script

#  Add your unittests of "script" module here

You will need to make sure that your main script logic is wrapped inside an if __name__ == "__main__": block, however. Usually I’ll just put the main script code in a “main()” function and call it from the “if” block:

def main():
    #  script code here

if __name__ == "__main__":
    main()

Another way of importing a script as a module is:

import sys
from pathlib import Path
import types

script_full_path = Path("script").resolve()

with open(script_full_path, "r") as fp:
    module_bytecode = compile(fp.read(), script_full_path.parent, 'exec')
script = types.ModuleType("script")
exec(module_bytecode, module.__dict__)
sys.modules["script"] = script

#  Add your unittests of "script" module here