Passing cython function pointers through python

atest.py

import pyximport
import numpy as np
pyximport.install()
cmap = pyximport.load_module('cmap', 'cmap.pyx')
cadd = pyximport.load_module('cadd', 'cadd.pyx')

a = np.arange(6, dtype='i')
print('start', a)
cmap.pymap(a, cadd.pyadd_one)
print('add_one', a)
r = cmap.pyreduce(a, cadd.pyadd_two)
print('reduce', r)

cadd.pyx

from cmap cimport WrapMap, WrapReduce

cdef int add_one(int x) nogil:
  return x + 1

cdef int add_two(int x, int y) nogil:
  return x + y

pyadd_one = WrapMap(add_one)

pyadd_two = WrapReduce(add_two)

cmap.pxd

ctypedef int (*mymap_t)(int x) nogil
ctypedef int (*myreduce_t)(int x, int y) nogil

cdef class PyMap(object):
  cdef mymap_t func
  cdef void _set(self, mymap_t func)

cdef PyMap WrapMap(mymap_t func)

cdef class PyReduce(object):
  cdef myreduce_t func
  cdef void _set(self, myreduce_t func)

cdef PyReduce WrapReduce(myreduce_t func)

cdef void mymap(int[:] a, mymap_t func) nogil

cdef int myreduce(int[:] a, myreduce_t func) nogil

cpdef pymap(int[:] a, PyMap func)

cpdef int pyreduce(int[:] a, PyReduce func)

cmap.pyx

cimport cython

cdef class PyMap(object):
  cdef void _set(self, mymap_t func):
    self.func = func

cdef PyMap WrapMap(mymap_t func):
  cdef PyMap o = PyMap()
  o._set(func)
  return o

cdef class PyReduce(object):
  cdef void _set(self, myreduce_t func):
    self.func = func

cdef PyReduce WrapReduce(myreduce_t func):
  cdef PyReduce o = PyReduce()
  o._set(func)
  return o

@cython.boundscheck(False)
cdef void mymap(int[:] a, mymap_t func) nogil:
  cdef int i
  for i in range(a.shape[0]):
    a[i] = func(a[i])

@cython.boundscheck(False)
cdef int myreduce(int[:] a, myreduce_t func) nogil:
  if a.shape[0] < 0:
    return 0  # Really should be exception instead
  elif a.shape[0] < 2:
    return a[0]
  cdef int x = func(a[0], a[1])
  cdef int i
  for i in range(2, a.shape[0]):
    x = func(x, a[i])
  return x

cpdef pymap(int[:] a, PyMap func):
  mymap(a, func.func)

cpdef int pyreduce(int[:] a, PyReduce func):
  return myreduce(a, func.func)