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++"

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

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

@register_cell_magic
def cpp(line, cell, *what):
    """Compile, execute C++ code wrapped with pybind11 and import it"""
    args = line.split(" ")
    rebuild = False
    cmake_path = None
    function = ""
    run_cmake = False
    # TODO: Add full argument parser
    for i,arg in enumerate(args):
        if arg == "-c":
            cmake_path = args[i+1].strip()
        if arg == "-f":
            function = args[i+1].strip()
        if arg == "-rebuild":
            rebuild = True


    # We first retrieve the current IPython interpreter
    # instance.
    ip = get_ipython()

    cmake = ""
    if cmake_path is not None:
        with open(os.path.join(cmake_path,"CMakeLists.txt"), 'r') as f:
            cmake = f.read()
    # name = ''.join(random.choices(string.ascii_lowercase, k=10))
    hash_object = hashlib.sha256((cmake+" ".join(line)+function+cell).encode())
    name = "cpp_magic_"+hash_object.hexdigest()[:10]
    if cmake_path is None:
        ensure_tmp_folder()
    else:
        run_cmake = not ensure_build_dir(cmake_path, clear=rebuild)
        p = os.path.abspath(os.path.join(cmake_path,"build"))
        if p not in sys.path:
            sys.path.append(p)
        ensure_pybind(cmake_path)

    if cmake_path is not None:
        path = cmake_path
    else:
        path = tmp_path

    # We define the source and executable filenames.
    try:
        try:
            if rebuild:
                raise Exception()
            mdl = importlib.import_module(name)
        except:
            exe = sys.executable
            split = cell.split("defs")
            if len(split) == 1 and function == "":
                raise Exception("You have to name a function as argument ('%%cpp -f <functionname>') or manual set defs")
            split.append("")

            curr_defs = split[1]
            if function != "":
                curr_defs += def_template.format(function=function)


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


            if cmake_path is None:
                # We compile the C++ code into an executable.
                command = f"{compiler_cmd} {' '.join(cmd_flags)} -O3 -Wall -shared -std=c++11 -fPIC $({exe} -m pybind11 --includes) {tmp_path}/{name}.cpp -o {tmp_path}/{name}.so"
                compile = ip.getoutput(command)
                if len(compile) != 0:
                    raise Exception("\n".join(compile))
            else:
                with open(os.path.join(cmake_path,"CMakeLists.txt"), 'w') as f:
                    f.write(cmake.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(cmake_path, "build"))
                    if err:
                        raise Exception(output)
                    print(output)
                command = ["make"]
                print("make:")
                err, output = exec_command(command, cwd=os.path.join(cmake_path, "build"))
                if err:
                    raise Exception(output)
                print(output)
                with open(os.path.join(cmake_path,"CMakeLists.txt"), 'w') as f:
                    f.write(cmake)


        # We execute the executable and return the output.
        # get a handle on the module
        mdl = importlib.import_module(name)

        if os.path.exists(os.path.join(path,f"{name}.cpp")):
            os.remove(os.path.join(path,f"{name}.cpp"))

        # 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("_")]
    except Exception as e:
        with open(os.path.join(cmake_path,"CMakeLists.txt"), 'w') as f:
            f.write(cmake)
        if os.path.exists(os.path.join(path,f"{name}.cpp")):
            os.remove(os.path.join(path,f"{name}.cpp"))
        raise e

    # now drag them in
    ip.push({k: getattr(mdl, k) for k in names})