pR0Ps/inline_c.py

## 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)
	#!/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)