call.py

#!/usr/bin/env python

"""
Imports and calls a python function

Usage:
    call.py 'my.package:callable' arg1 arg2
"""

import argparse
import os
import sys
import importlib
import subprocess

from typing import List

def python_env_sys_path(env_prefix: str) -> List[str]:
    """
    Gets the relevant sys.path entries for a python environment
    by shelling out to its python and extracting sys.path.

    :param env_prefix:
        Path to python environment.  bin/python must exist here.

    :return:
        List of sys.path entries for the conda env.
    """
    command = [
        os.path.join(env_prefix, 'bin', 'python'),
        '-c',
        f"import sys; print('\\n'.join([p for p in sys.path if p.startswith('{env_prefix}')]))"
    ]
    return subprocess.check_output(command).decode('utf-8').strip().split('\n')


def include_python_env_sys_path(other_python_env_prefix: str):
    """
    Adds an external conda environment's sys.path entries
    to the current running python sys.path.
    Doing this lets you import python modules from the other
    conda environment, essentially 'stacking' this python
    on top of it.

    :param conda_env_prefix:
        Path to other conda environment.  bin/python must exist here.

    :return:
        List of path entries added to our sys.path
    """
    other_python_env_paths = python_env_sys_path(other_python_env_prefix)
    paths_to_add = [p for p in other_python_env_paths if p not in sys.path]
    sys.path += paths_to_add
    return paths_to_add


def import_it(name: str):
    """
    Imports a python module / symbol from a string name and returns it.
    It is up to you to make sure that the module is importable
    via python sys.path.

    :param name:
        Python module / symbol name.  Examples:
            * 'mypackage.thing' -> from mypackage import thing
            * 'mypackage:callable' -> from mypackage import callable
            * 'mypackage' -> import mypackage

    :return: imported Python module symbol
    """
    if ':' in name:
        module_name, symbol_name = name.rsplit(':', 1)
    elif '.' in name:
        module_name, symbol_name = name.rsplit('.', 1)
    else:
        module_name = name
        symbol_name = None

    module = importlib.import_module(module_name)
    if symbol_name is None:
        return module
    return getattr(module, symbol_name)


def call(name: str, *args):
    """
    Calls a python function by string name.

    :param name:
        Python module callable name.
        E.g. my.package:callable

    :param *args: args to pass to callable.

    :return: result of callable
    """
    return import_it(name)(*args)


def main(argv = sys.argv[1:]):
    """
    Imports callable and calls it with args.

    If --base_env is given as a comma separated list of prefix paths to a python env
    (e.g. a conda environemnt prefix path), sys.path will be altered
    to include that environment's sys.path as well. This allows
    us to call a function that might have dependencies in another environment,
    effectively 'stacking' the current python env on the base env(s).
    """

    parser = argparse.ArgumentParser(description=main.__doc__)
    parser.add_argument('--base_env', action='append', dest='base_env', type=str, help='Python environment prefix to include in sys.path.')
    parser.add_argument('callable')
    parser.add_argument('args', nargs='*')
    args = parser.parse_args(argv)

    print(args)
    if args.base_env is not None:
        # base_envs = args.base_env.split(',')
        for prefix in args.base_env:
            include_python_env_sys_path(prefix)
    call(args.callable, *args.args)


if __name__ == '__main__':
    main()