pybindmagic/pybindmagic/__init__.py

from IPython.core.magic import register_cell_magic, needs_local_scope

import importlib
import glob
import sys, os
import random
import string
import hashlib
import shutil
import subprocess

compiler_cmd = "c++"

optimize = "-O0"

cmd_flags = ["$(pkg-config --cflags eigen3)", "-Wall"]

template = """
#include <pybind11/pybind11.h>
namespace py = pybind11;

{cell_code}

PYBIND11_MODULE({name}, m)
    {defs}
"""

def_template = '\nm.def("{function}", &{function});\n'

TMP_PATH = "pbm_modules"

# Make sure the correct python executeable is used using cmake
os.environ["PATH"] = os.path.dirname(sys.executable) + os.pathsep + os.environ["PATH"]

def exec_command(cmd, cwd=None,**kwargs):
    try:
        output = subprocess.check_output(cmd,cwd=cwd,stderr=subprocess.STDOUT,**kwargs)
        if isinstance(output,bytes):
            return False, output.decode("utf-8")
        return False, output
    except subprocess.CalledProcessError as e:
        if isinstance(e.output,bytes):
            return True, e.output.decode("utf-8")
        return True, e.output


def ensure_tmp_folder():
    path = os.path.join(os.getcwd(), TMP_PATH)
    if not os.path.exists(path):
        os.makedirs(path)
    if not path in sys.path:
        sys.path.append(path)

def ensure_pybind(path, force=False):
    path = os.path.join(path,"pybind11")
    exists = os.path.exists(path)
    if exists and force:
        shutil.rmtree(os.path.join(path,"pybind11"))
    if not exists:
        err, output = exec_command(["git", "clone","https://github.com/pybind/pybind11.git",path])

def ensure_build_dir(path, clear=False):
    path = os.path.join(path, "build")
    exists = os.path.exists(path)
    if exists and clear:
        shutil.rmtree(path)
    if not exists or clear:
        os.makedirs(path)
    return exists and not clear

def parse_args(line):
    args = line.split(" ")
    ret = {}
    ret["cmake_path"] = None
    ret["functions"] = []
    ret["rebuild"] = False
    # TODO: Add full argument parser
    for i,arg in enumerate(args):
        if arg == "-c":
            ret["cmake_path"] = args[i+1].strip()
        if arg == "-f":
            ret["functions"].append(args[i+1].strip())
        if arg == "-rebuild":
            ret["rebuild"] = True
    return ret


def get_CMakeLists(args):
    CMakeLists = ""
    if args["cmake_path"] is not None:
        with open(os.path.join(args["cmake_path"],"CMakeLists.txt"), 'r') as f:
            CMakeLists = f.read()
    return CMakeLists

def cleanup(args, name, CMakeLists):
    if args["cmake_path"] is not None:
        with open(os.path.join(args["cmake_path"],"CMakeLists.txt"), 'w') as f:
            f.write(CMakeLists)
    if os.path.exists(os.path.join(args["path"],f"{name}.cpp")):
        os.remove(os.path.join(args["path"],f"{name}.cpp"))

def get_gcc_cmd(name):
    pyexe = sys.executable
    return f"{compiler_cmd} {' '.join(cmd_flags)} {optimize} -std=c++11 -fPIC $({pyexe} -m pybind11 --includes) -c {TMP_PATH}/{name}.cpp -o {TMP_PATH}/{name}.o"

def get_linker_cmd(name):
    others = " ".join([f"-l{os.path.basename(n)[3:-3]}" for n in sorted(glob.glob(f"{TMP_PATH}/*.so"), key=os.path.getmtime)[::-1]])
    return f"{compiler_cmd} -Wl,-rpath {TMP_PATH} -shared {TMP_PATH}/{name}.o -o {TMP_PATH}/{name}.so -L{TMP_PATH} {others}"

def get_unique_linker_cmd(other_path, rnd_path):
    # TODO: Merge with get_linker_cmd
    other = os.path.basename(other_path)[:-2]

    all_other = " ".join([f"-l{os.path.basename(n)[3:-3]}" for n in sorted(glob.glob(f"{TMP_PATH}/*.so"), key=os.path.getmtime)[::-1] if os.path.basename(n)[3:-3] != os.path.basename(other_path)[3:-2]])
    return f"{compiler_cmd} -Wl,-rpath {TMP_PATH} -shared {TMP_PATH}/{other}.o -o {rnd_path}/{other}.so -L{TMP_PATH} {all_other}"


def import_to_ip(name):
    # get a handle on the module
    mdl = importlib.import_module(name)
    # is there an __all__?  if so respect it
    if "__all__" in mdl.__dict__:
        names = mdl.__dict__["__all__"]
    else:
        # otherwise we import all names that don't begin with _
        names = [x for x in mdl.__dict__ if not x.startswith("_")]
    get_ipython().push({k: getattr(mdl, k) for k in names})

@register_cell_magic
def cpp(line, cell, *what):
    """Compile, execute C++ code wrapped with pybind11 and import it"""
    args = parse_args(line)
    ip = get_ipython()
    # get name of module
    if args["cmake_path"] is None:
        # TODO!
        compile_id = get_gcc_cmd("dummy")
    else:
        compile_id = get_CMakeLists(args)

    hash_object = hashlib.sha256((compile_id+"|"+line+"|"+cell).encode())
    name = "libcpp_magic_"+hash_object.hexdigest()[:10]
    run_cmake = False

    # init temp folder and pybind etc
    if args["cmake_path"] is None:
        ensure_tmp_folder()
        args["path"] = TMP_PATH
        CMakeLists = None
    else:
        ensure_build_dir(args["cmake_path"], clear=args["rebuild"])
        p = os.path.abspath(os.path.join(args["cmake_path"],"build"))
        if p not in sys.path:
            sys.path.append(p)
        run_cmake = not os.path.exists(os.path.join(p, "Makefile"))
        ensure_pybind(args["cmake_path"])
        args["path"] = args["cmake_path"]
        CMakeLists = compile_id
    build_needed = False
    try:
        if args["rebuild"]:
            raise Exception()
        mdl = importlib.import_module(name)
    except:
        build_needed = True

    if build_needed:
        try:
            # add pybind11 function defs
            split = cell.split("defs")
            if len(split) == 1 and len(args["functions"]) == 0:
                raise Exception("You have to name a function as argument ('%%cpp -f <functionname>') or manual set defs")
            # make sure that split is at least length 2
            split.append("")
            curr_defs = split[1]
            for function in args["functions"]:
                curr_defs += def_template.format(function=function)

            # We write the code to the C++ file.
            with open(os.path.join(args["path"],f"{name}.cpp"), 'w') as f:
                f.write(template.format(cell_code=split[0],name=name,defs="{"+curr_defs+"}"))

            if args["cmake_path"] is None:
                # We compile the C++ code into an object file
                command = get_gcc_cmd(name)
                compile = ip.getoutput(command)
                # currently warnings will lead to an abort
                if len(compile) != 0:
                    raise Exception("\n".join(compile))

                # link into a shared object
                linker = get_linker_cmd(name)
                link = ip.getoutput(linker)
                if len(link) != 0:
                    raise Exception("\n".join(link))


                # create unique folder
                rnd = ''.join(random.choice(string.ascii_uppercase + string.digits) for _ in range(5))
                rnd_path = os.path.join(TMP_PATH, rnd)
                if not os.path.exists(rnd_path):
                    os.makedirs(rnd_path)
                    with open(os.path.join(rnd_path,"__init__.py"), 'w+') as f:
                        pass
                # Rethink design decision: Maybe do not import all others again, maybe give cells names and reduce the amount of linking
                # update all others by relinking and reimporting orderd by compilation time
                for other_path in sorted(glob.glob(f"{TMP_PATH}/*.o"), key=os.path.getmtime):
                    other = os.path.basename(other_path)[:-2]


                    linker = get_unique_linker_cmd(other_path, rnd_path)
                    link = ip.getoutput(linker)
                    if len(link) != 0:
                        raise Exception("\n".join(link))

                    import_to_ip(f"{rnd}.{other}")

                shutil.rmtree(rnd_path)

            else:
                with open(os.path.join(args["cmake_path"],"CMakeLists.txt"), 'w') as f:
                    f.write(CMakeLists.replace("{name}",name))
                # TODO: Add option to add cmake flags
                if run_cmake:
                    command = ["cmake", ".."]
                    print("cmake:")
                    err, output = exec_command(command, cwd=os.path.join(args["cmake_path"], "build"))
                    if err:
                        raise Exception(output)
                    print(output)
                command = ["make"]
                print("make:")
                err, output = exec_command(command, cwd=os.path.join(args["cmake_path"], "build"))
                if err:
                    raise Exception(output)
                print(output)
                with open(os.path.join(args["cmake_path"],"CMakeLists.txt"), 'w') as f:
                    f.write(CMakeLists)
        except Exception as e:
            # clean up even if error occurs
            cleanup(args, name, CMakeLists)
            raise e

    import_to_ip(name)

    # clean up
    cleanup(args, name, CMakeLists)