pythonimport-libraries

Replacing imp with importlib maintaining old behavior?


I inherited some code that I need to rework since it is using deprecated imp module instead of importlib

To test the functionality I have created a simple test script:

# test.py
def main():
    print("main()")


if __name__ == "__main__":
    print("__main__")
    main()

When I run that with the old code (below as a minimal example)

# original imp based minimal example
import imp
import os
import site

script = r"./test.py"
script_dir = os.path.dirname(script)
script_name = os.path.splitext(os.path.basename(script))[0]

site.addsitedir(script_dir)

fp, pathname, description = imp.find_module(script_name, [script_dir])
imp.load_source('__main__', pathname, fp)

Would output the following:

__main__
main()

Now I have some trouble mimicking this verbatim using the importlib module since I do not get the output as above with the following code:

# my try so far using importlib
import importlib
import os
import site

script = os.path.abspath(r"./test.py")
script_dir = os.path.dirname(script)
script_name = os.path.splitext(os.path.basename(script))[0]

site.addsitedir(script_dir)

spec = importlib.util.spec_from_file_location(script_name, script)
module = importlib.util.module_from_spec(spec)
sys.modules[script_name] = module
spec.loader.exec_module(module)

Instead it outputs nothing. :-( Any suggestions on how to mimic the old (imp) behavior using importlib?


Solution

  • Eventually I solved it like below based on https://github.com/epfl-scitas/spack/blob/a60ae07083a5744607064221a0cd48204a54394e/lib/spack/llnl/util/lang.py#L598-L625:

    if sys.version_info[0] == 3: 
        if sys.version_info[1] >= 5:
            import importlib.util
            spec = importlib.util.spec_from_file_location(module_name, module_path)
            module = importlib.util.module_from_spec(spec)
            spec.loader.exec_module(module)
        elif sys.version_info[1] < 5:
            import importlib.machinery
            loader = importlib.machinery.SourceFileLoader(module_name, module_path)
            loader.load_module()
    elif sys.version_info[0] == 2:
        import imp
        imp.load_source(module_name, module_path)