Use inline C code in Python (using Cython)

inline_c.py

#!/usr/bin/env python

import cython
import functools
import inspect
import textwrap


_ARRAY_NAME = "Array_{c_type}"
_MAKE_ARRAY = _ARRAY_NAME + "(&{name}[0], {name}.shape[0])"
_PYTHON_TYPEDEF = "ctypedef struct " + _ARRAY_NAME + ":\n    {c_type} *data\n    unsigned int size"
_C_TYPEDEF = "typedef struct " + _ARRAY_NAME + " {{\n  {c_type} *data;\n  unsigned int size;\n}} " + _ARRAY_NAME + ";\n"
_FUNCTION_DEFINITION = "{return_type} {function_name}({c_params})"
_FUNCTION_TEMPLATE = "{function_definition}{{\n{function_body}\n}}"
_CYTHON_TEMPLATE = '''
cdef extern from *:
    """
{c_typedefs}
{function_code}
    """
{python_typedefs}
    {function_definition}

def __stub({orig_c_params}):
    return {function_name}({call_params})
'''

def inline_c(func):
    """Compiles C code from a docstring and returns a callable function for it.

    Type annotations must be used to communicate the parameter and return
    types. Currently only simple types that are automatically mapped by Cython
    are supported.

    The returned function will have the docstring set to the C code that was
    used to generate it for better introspectability.

    Ex:
    ```
    >>> @inline_c
    ... def fibonacci(n: "unsigned int") -> "unsigned int":
    ...     '''
    ...     if (n < 2){
    ...         return n;
    ...     }
    ...     return fibonacci(n - 1) + fibonacci(n - 2);
    ...     '''
    >>> print(fibonacci.__doc__)
    Compiled from the following code:
    ```
    unsigned int fibonacci(unsigned int n){
      if (n < 2){
          return n;
      }
      return fibonacci(n - 1) + fibonacci(n - 2);
    }
    ```
    >>> fibonacci(20)
    6765
    ```

    Lists are supported, but require using `array` objects from the `array`
    module in the stdlib. To use lists, annotate the type in the same way that
    a Cython memoryview would be typed (`type[:]`). To make these arrays
    available in C code, they will automatically be converted into a struct
    with `data` and `size` attributes.

    Currently return types cannot be lists. A workaround is to modify the data
    in-place or pass in a pre-created array that the C function can load data
    into.

    Currently only single-dimension memoryviews are supported.

    Ex:
    ```
    >>> @inline_c
    ... def func(l: "int[:]") -> "void":
    ...    r'''
    ...    printf("first: %d\\nlast: %d\\nlen: %d\\n", l.data[0], l.data[l.size-1], l.size);
    ...    '''
    >>> import array
    >>> func(array.array("i", [1, 10, 100, 999]))
    first: 1
    last: 999
    len: 4
    """

    if func.__code__.co_flags & inspect.CO_VARARGS:
        raise ValueError("Compiled functions cannot use varags (*args)")
    if func.__code__.co_flags & inspect.CO_VARKEYWORDS:
        raise ValueError("Compiled functions cannot use keyword args (**kwargs)")

    function_name = func.__name__

    doc = inspect.getdoc(func)
    if not doc:
        raise ValueError("No code to compile supplied in docstring")
    function_body = textwrap.indent(doc, "  ")

    missing = set(func.__code__.co_varnames[:func.__code__.co_argcount])
    params = []
    return_type = None
    for k, v in func.__annotations__.items():
        if not isinstance(v, str):
            raise ValueError(
                "Type annotation for '{}' is required to be a string contianing a c type, not {}".format(k, type(v).__name__)
            )
        if k == "return":
            return_type = v
        else:
            missing.remove(k)
            params.append((k, v, v.rstrip(" :[]") if v.endswith("[:]") else None))

    if not return_type:
        raise ValueError("Return type not specified")
    if missing:
        raise ValueError("The following args were not type-annotated: {}".format(", ".join(missing)))

    orig_c_params = ", ".join("{1} {0}".format(*x) for x in params)
    c_params = ", ".join("{} {}".format(_ARRAY_NAME.format(c_type=x[2]) if x[2] else x[1], x[0]) for x in params)
    call_params = ", ".join(_MAKE_ARRAY.format(c_type=x[2], name=x[0]) if x[2] else x[0] for x in params)
    function_definition = _FUNCTION_DEFINITION.format(**locals())

    # Make the required typedefs. Normally it would be enough to use `ctypedef
    # public struct Thing` in Cython code to have the struct be available in
    # both the C and Cython scope, but due to issues with Cython's regex code
    # parsing, it's not possible to use a public typedef using cython.inline.
    # The workaround here is to generate code for identical Cython and C
    # typedefs and add them to both sections.
    # Additionally, the difference between a blank line and one with just
    # whitespace in Cython matters to the parser.
    required_typedefs = set(x[2] for x in params if x[2])
    python_typedefs = textwrap.indent("\n".join(_PYTHON_TYPEDEF.format(c_type=x) for x in required_typedefs), "    ") or "    "
    c_typedefs = textwrap.indent("\n".join(_C_TYPEDEF.format(c_type=x) for x in required_typedefs), "    ")

    function_code = textwrap.indent(_FUNCTION_TEMPLATE.format(**locals()), "    ")
    cython_code=_CYTHON_TEMPLATE.format(**locals())
    try:
        newfunc = cython.inline(cython_code, quiet=True)["__stub"]
    except Exception as e:
        raise ValueError("Failed to compile the following code:\n\n{}".format(cython_code)) from e

    functools.update_wrapper(newfunc, func)
    newfunc.__doc__ = "Compiled from the following code:\n```\n{}\n```".format(_FUNCTION_TEMPLATE.format(**locals()))

    return newfunc


####################
# Examples / testing
####################

if __name__ == "__main__":

    import array
    import contextlib

    def benchmark(code, num=5000000):
        import timeit
        print("{:<30} {:>11}: ".format(code, "(x{})".format(num)), end="", flush=True)
        print("{:23.20f}".format(timeit.timeit(code, number=num, globals=globals())))

    @contextlib.contextmanager
    def section(name):
        print(name)
        print("-" * len(name))
        try:
            yield
        finally:
            print("\n")


    with section("Factorial"):
        def py_factorial(n):
            t = 1
            for i in range(1, n+1):
                t = t*i
            return t

        @inline_c
        def c_factorial(n: "unsigned int") -> "unsigned long":
            """
            unsigned long t = 1;
            for (unsigned int i=1; i < n+1; i++){
                t *= i;
            }
            return t;
            """

        assert py_factorial(20) == c_factorial(20)
        benchmark("py_factorial(20)")
        benchmark("c_factorial(20)")

    with section("Multiply"):
        def py_multiply(a, b):
            return a*b

        @inline_c
        def c_multiply(a: "long", b: "long") -> "long":
            """
            return a*b;
            """

        assert py_multiply(123789, 12912) == c_multiply(123789, 12912)
        benchmark("py_multiply(123789, 12912)", num=50000000)
        benchmark("c_multiply(123789, 12912)", num=50000000)

    with section("Sum"):
        @inline_c
        def c_sum(data: "long[:]") -> "long":
            """
            long t = 0;
            for (unsigned int i = 0; i < data.size; i++) {
                t += data.data[i];
            }
            return t;
            """

        def py_sum_opt(data):
            return sum(data)

        def py_sum(data):
            t = 0;
            for x in data:
                t += x
            return t

        data = array.array("l", range(50))
        assert py_sum(data) == py_sum_opt(data) == c_sum(data)
        benchmark("py_sum(data)")
        benchmark("py_sum_opt(data)")
        benchmark("c_sum(data)")

    with section("Invert list"):
        @inline_c
        def c_invert_list(data: "long[:]") -> "void":
            """
            for (unsigned int i = 0; i < data.size; i++) {
                data.data[i] = -data.data[i];
            }
            """

        def py_invert_list(data):
            for i, x in enumerate(data):
                data[i] = -x

        data1 = array.array("l", range(50))
        data2 = array.array("l", range(50))
        c_invert_list(data1)
        py_invert_list(data2)
        assert data1 == data2
        benchmark("py_invert_list(data1)")
        benchmark("c_invert_list(data1)")

    with section("Fibonacci"):
        @inline_c
        def c_fibonacci(n: "unsigned int") -> "unsigned int":
            """
            if (n < 2){
                return n;
            }
            return c_fibonacci(n - 1) + c_fibonacci(n - 2);
            """

        @inline_c
        def c_fibonacci_opt(n: "unsigned int") -> "unsigned int":
            """
            if (n < 2){
                return n;
            }
            unsigned int buf[3] = {0, 1, 1};
            for (unsigned int i = 3; i <= n; i++){
                buf[0] = buf[1];
                buf[1] = buf[2];
                buf[2] = buf[0] + buf[1];
            }
            return buf[2];
            """

        def py_fibonacci(n):
            if n < 2:
                return n;
            return py_fibonacci(n - 1) + py_fibonacci(n - 2)

        def py_fibonacci_opt(n):
            if n < 2:
                return n;
            buf = [0, 1, 1]
            for x in range(3, n + 1):
                buf[0] = buf[1]
                buf[1] = buf[2]
                buf[2] = buf[0] + buf[1]

            return buf[2]

        assert py_fibonacci(10) == py_fibonacci_opt(10) == c_fibonacci(10) == c_fibonacci_opt(10)
        benchmark("py_fibonacci(30)", num=100)
        benchmark("py_fibonacci_opt(30)", num=100)
        benchmark("c_fibonacci(30)", num=100)
        benchmark("c_fibonacci_opt(30)", num=100)

Results on a 2015 MBP (2.7GHz i7):

$ ./inline_c.py 
Factorial
---------
py_factorial(20)                (x5000000):  6.00467820093035697937
c_factorial(20)                 (x5000000):  0.39383666706271469593


Multiply
--------
py_multiply(123789, 12912)     (x50000000):  5.39791185897774994373
c_multiply(123789, 12912)      (x50000000):  4.38778426009230315685


Sum
---
py_sum(data)                    (x5000000): 10.62241740315221250057
py_sum_opt(data)                (x5000000):  2.88693389389663934708
c_sum(data)                     (x5000000):  1.25900624203495681286


Invert list
-----------
py_invert_list(data1)           (x5000000): 24.67182233510538935661
c_invert_list(data1)            (x5000000):  1.20100786117836833000


Fibonacci
---------
py_fibonacci(30)                    (x100): 27.24636415718123316765
py_fibonacci_opt(30)                (x100):  0.00051718694157898426
c_fibonacci(30)                     (x100):  0.39685972803272306919
c_fibonacci_opt(30)                 (x100):  0.00000962615013122559