# High-level extension API

This extension API is exposed through the `numba.extending` module.

## Implementing functions

The `@overload` decorator allows you to implement arbitrary functions for use in nopython mode functions. The function decorated with `@overload` is called at compile-time with the types of the function’s runtime arguments. It should return a callable representing the implementation of the function for the given types. The returned implementation is compiled by Numba as if it were a normal function decorated with `@jit`. Additional options to `@jit` can be passed as dictionary using the `jit_options` argument.

For example, let’s pretend Numba doesn’t support the `len()` function on tuples yet. Here is how to implement it using `@overload`:

    from numba import types
    from numba.extending import overload

    @overload(len)
    def tuple_len(seq):
       if isinstance(seq, types.BaseTuple):
           n = len(seq)
           def len_impl(seq):
               return n
           return len_impl

You might wonder, what happens if `len()` is called with something else than a tuple? If a function decorated with `@overload` doesn’t return anything (i.e. returns None), other definitions are tried until one succeeds. Therefore, multiple libraries may overload `len()` for different types without conflicting with each other.

## Implementing methods

The `@overload_method` decorator similarly allows implementing a method on a type well-known to Numba.

numba.core.extending.overload_method(typ, attr, **kwargs)

  A decorator marking the decorated function as typing and implementing method attr for the given Numba type in nopython mode.

  kwargs are passed to the underlying @overload call.

  Here is an example implementing .take() for array types:

      @overload_method(types.Array, 'take')
      def array_take(arr, indices):
          if isinstance(indices, types.Array):
              def take_impl(arr, indices):
                  n = indices.shape[0]
                  res = np.empty(n, arr.dtype)
                  for i in range(n):
                      res[i] = arr[indices[i]]
                  return res
              return take_impl

## Implementing classmethods

The `@overload_classmethod` decorator similarly allows implementing a classmethod on a type well-known to Numba.

numba.core.extending.overload_classmethod(typ, attr, **kwargs)

  A decorator marking the decorated function as typing and implementing classmethod attr for the given Numba type in nopython mode.

  Similar to `overload_method`.

  Here is an example implementing a classmethod on the Array type to call `np.arange()`:

      @overload_classmethod(types.Array, "make")
      def ov_make(cls, nitems):
          def impl(cls, nitems):
              return np.arange(nitems)
          return impl

  The above code will allow the following to work in jit-compiled code:

      @njit
      def foo(n):
          return types.Array.make(n)

## Implementing attributes

The `@overload_attribute` decorator allows implementing a data attribute (or property) on a type. Only reading the attribute is possible; writable attributes are only supported through the low-level API.

The following example implements the `nbytes` attribute on Numpy arrays:

    @overload_attribute(types.Array, 'nbytes')
    def array_nbytes(arr):
       def get(arr):
           return arr.size * arr.itemsize
       return get

## Importing Cython Functions

The function `get_cython_function_address` obtains the address of a C function in a Cython extension module. The address can be used to access the C function via a `ctypes.CFUNCTYPE()` callback, thus allowing use of the C function inside a Numba jitted function. For example, suppose that you have the file `foo.pyx`:

    from libc.math cimport exp

    cdef api double myexp(double x):
        return exp(x)

You can access `myexp` from Numba in the following way:

    import ctypes
    from numba.extending import get_cython_function_address

    addr = get_cython_function_address("foo", "myexp")
    functype = ctypes.CFUNCTYPE(ctypes.c_double, ctypes.c_double)
    myexp = functype(addr)

The function `myexp` can now be used inside jitted functions, for example:

    @njit
    def double_myexp(x):
        return 2*myexp(x)

One caveat is that if your function uses Cython’s fused types, then the function’s name will be mangled. To find out the mangled name of your function you can check the extension module’s `__pyx_capi__` attribute.

## Implementing intrinsics

The `@intrinsic` decorator is used for marking a function func as typing and implementing the function in `nopython` mode using the llvmlite IRBuilder API. This is an escape hatch for expert users to build custom LLVM IR that will be inlined into the caller, there is no safety net!

The first argument to func is the typing context. The rest of the arguments corresponds to the type of arguments of the decorated function. These arguments are also used as the formal argument of the decorated function. If func has the signature `foo(typing_context, arg0, arg1)`, the decorated function will have the signature `foo(arg0, arg1)`.

The return values of func should be a 2-tuple of expected type signature, and a code-generation function that will passed to `lower_builtin()`. For an unsupported operation, return `None`.

Here is an example that cast any integer to a byte pointer:

    from numba import types
    from numba.extending import intrinsic

    @intrinsic
    def cast_int_to_byte_ptr(typingctx, src):
        # check for accepted types
        if isinstance(src, types.Integer):
            # create the expected type signature
            result_type = types.CPointer(types.uint8)
            sig = result_type(types.uintp)
            # defines the custom code generation
            def codegen(context, builder, signature, args):
                # llvm IRBuilder code here
                [src] = args
                rtype = signature.return_type
                llrtype = context.get_value_type(rtype)
                return builder.inttoptr(src, llrtype)
            return sig, codegen

it may be used as follows:

    from numba import njit

    @njit('void(int64)')
    def foo(x):
        y = cast_int_to_byte_ptr(x)

    foo.inspect_types()

and the output of `.inspect_types()` demonstrates the cast (note the `uint8*`):

    def foo(x):

        #   x = arg(0, name=x)  :: int64
        #   $0.1 = global(cast_int_to_byte_ptr: <intrinsic cast_int_to_byte_ptr>)  :: Function(<intrinsic cast_int_to_byte_ptr>)
        #   $0.3 = call $0.1(x, func=$0.1, args=[Var(x, check_intrin.py (24))], kws=(), vararg=None)  :: (uint64,) -> uint8*
        #   del x
        #   del $0.1
        #   y = $0.3  :: uint8*
        #   del y
        #   del $0.3
        #   $const0.4 = const(NoneType, None)  :: none
        #   $0.5 = cast(value=$const0.4)  :: none
        #   del $const0.4
        #   return $0.5

        y = cast_int_to_byte_ptr(x)

## Implementing mutable structures

Warning

This is an experimental feature, the API may change without warning.

The `numba.experimental.structref` module provides utilities for defining mutable pass-by-reference structures, a `StructRef`. The following example demonstrates how to define a basic mutable structure:

### Defining a StructRef

from `numba/tests/doc_examples/test_structref_usage.py`

     1import numpy as np
     2
     3from numba import njit
     4from numba.core import types
     5from numba.experimental import structref
     6
     7from numba.tests.support import skip_unless_scipy
     8
     9
    10# Define a StructRef.
    11# `structref.register` associates the type with the default data model.
    12# This will also install getters and setters to the fields of
    13# the StructRef.
    14@structref.register
    15class MyStructType(types.StructRef):
    16    def preprocess_fields(self, fields):
    17        # This method is called by the type constructor for additional
    18        # preprocessing on the fields.
    19        # Here, we don't want the struct to take Literal types.
    20        return tuple((name, types.unliteral(typ)) for name, typ in fields)
    21
    22
    23# Define a Python type that can be use as a proxy to the StructRef
    24# allocated inside Numba. Users can construct the StructRef via
    25# the constructor for this type in python code and jit-code.
    26class MyStruct(structref.StructRefProxy):
    27    def __new__(cls, name, vector):
    28        # Overriding the __new__ method is optional, doing so
    29        # allows Python code to use keyword arguments,
    30        # or add other customized behavior.
    31        # The default __new__ takes `*args`.
    32        # IMPORTANT: Users should not override __init__.
    33        return structref.StructRefProxy.__new__(cls, name, vector)
    34
    35    # By default, the proxy type does not reflect the attributes or
    36    # methods to the Python side. It is up to users to define
    37    # these. (This may be automated in the future.)
    38
    39    @property
    40    def name(self):
    41        # To access a field, we can define a function that simply
    42        # return the field in jit-code.
    43        # The definition of MyStruct_get_name is shown later.
    44        return MyStruct_get_name(self)
    45
    46    @property
    47    def vector(self):
    48        # The definition of MyStruct_get_vector is shown later.
    49        return MyStruct_get_vector(self)
    50
    51
    52@njit
    53def MyStruct_get_name(self):
    54    # In jit-code, the StructRef's attribute is exposed via
    55    # structref.register
    56    return self.name
    57
    58
    59@njit
    60def MyStruct_get_vector(self):
    61    return self.vector
    62
    63
    64# This associates the proxy with MyStructType for the given set of
    65# fields. Notice how we are not constraining the type of each field.
    66# Field types remain generic.
    67structref.define_proxy(MyStruct, MyStructType, ["name", "vector"])

The following demonstrates using the above mutable struct definition:

from `test_type_definition` of `numba/tests/doc_examples/test_structref_usage.py`

     1# Let's test our new StructRef.
     2
     3# Define one in Python
     4alice = MyStruct("Alice", vector=np.random.random(3))
     5
     6# Define one in jit-code
     7@njit
     8def make_bob():
     9    bob = MyStruct("unnamed", vector=np.zeros(3))
    10    # Mutate the attributes
    11    bob.name = "Bob"
    12    bob.vector = np.random.random(3)
    13    return bob
    14
    15bob = make_bob()
    16
    17# Out: Alice: [0.5488135  0.71518937 0.60276338]
    18print(f"{alice.name}: {alice.vector}")
    19# Out: Bob: [0.88325739 0.73527629 0.87746707]
    20print(f"{bob.name}: {bob.vector}")
    21
    22# Define a jit function to operate on the structs.
    23@njit
    24def distance(a, b):
    25    return np.linalg.norm(a.vector - b.vector)
    26
    27# Out: 0.4332647200356598
    28print(distance(alice, bob))

### Defining a method on StructRef

Methods and attributes can be attached using `@overload_*` as shown in the previous sections.

The following demonstrates the use of `@overload_method` to insert a method for instances of `MyStructType`:

from `test_overload_method` of `numba/tests/doc_examples/test_structref_usage.py`

     1from numba.core.extending import overload_method
     2from numba.core.errors import TypingError
     3
     4# Use @overload_method to add a method for
     5# MyStructType.distance(other)
     6# where *other* is an instance of MyStructType.
     7@overload_method(MyStructType, "distance")
     8def ol_distance(self, other):
     9    # Guard that *other* is an instance of MyStructType
    10    if not isinstance(other, MyStructType):
    11        raise TypingError(
    12            f"*other* must be a {MyStructType}; got {other}"
    13        )
    14
    15    def impl(self, other):
    16        return np.linalg.norm(self.vector - other.vector)
    17
    18    return impl
    19
    20# Test
    21@njit
    22def test():
    23    alice = MyStruct("Alice", vector=np.random.random(3))
    24    bob = MyStruct("Bob", vector=np.random.random(3))
    25    # Use the method
    26    return alice.distance(bob)

### `numba.experimental.structref` API Reference

Utilities for defining a mutable struct.

A mutable struct is passed by reference; hence, structref (a reference to a struct).

class numba.experimental.structref.StructRefProxy(*args)

  A PyObject proxy to the Numba allocated structref data structure.

  Notes

  - Subclasses should not define `__init__`.

  - Subclasses can override `__new__`.

numba.experimental.structref.define_attributes(struct_typeclass)

  Define attributes on struct_typeclass.

  Defines both setters and getters in jit-code.

  This is called directly in register().

numba.experimental.structref.define_boxing(struct_type, obj_class)

  Define the boxing & unboxing logic for struct_type to obj_class.

  Defines both boxing and unboxing.

  - boxing turns an instance of struct_type into a PyObject of obj_class

  - unboxing turns an instance of obj_class into an instance of struct_type in jit-code.

  Use this directly instead of define_proxy() when the user does not want any constructor to be defined.

numba.experimental.structref.define_constructor(py_class, struct_typeclass, fields)

  Define the jit-code constructor for struct_typeclass using the Python type py_class and the required fields.

  Use this instead of define_proxy() if the user does not want boxing logic defined.

numba.experimental.structref.define_proxy(py_class, struct_typeclass, fields)

  Defines a PyObject proxy for a structref.

  This makes py_class a valid constructor for creating a instance of struct_typeclass that contains the members as defined by fields.

  Parameters:
    py_class type

      The Python class for constructing an instance of struct_typeclass.

    struct_typeclass numba.core.types.Type

      The structref type class to bind to.

    fields Sequence[str]

      A sequence of field names.

  Returns:
    None

numba.experimental.structref.register(struct_type)

  Register a numba.core.types.StructRef for use in jit-code.

  This defines the data-model for lowering an instance of struct_type. This defines attributes accessor and mutator for an instance of struct_type.

  Parameters:
    struct_type type

      A subclass of numba.core.types.StructRef.

  Returns:
    struct_type type

      Returns the input argument so this can act like a decorator.

  Examples

      class MyStruct(numba.core.types.StructRef):
          ...  # the simplest subclass can be empty

      numba.experimental.structref.register(MyStruct)

numba.experimental.structref.structref_is(context, builder, sig, args)

  Define the ‘is’ operator for structrefs by comparing the memory addresses. This is the identity check for structref objects.

## Determining if a function is already wrapped by a `jit` family decorator

The following function is provided for this purpose.

extending.is_jitted()

  Returns True if a function is wrapped by one of the Numba @jit decorators, for example: numba.jit, numba.njit

  The purpose of this function is to provide a means to check if a function is already JIT decorated.