# Compiling Python classes with `@jitclass`

Note

This is a early version of jitclass support. Not all compiling features are exposed or implemented, yet.

Numba supports code generation for classes via the `numba.experimental.jitclass()` decorator. A class can be marked for optimization using this decorator along with a specification of the types of each field. We call the resulting class object a jitclass. All methods of a jitclass are compiled into nopython functions. The data of a jitclass instance is allocated on the heap as a C-compatible structure so that any compiled functions can have direct access to the underlying data, bypassing the interpreter.

## Basic usage

Here’s an example of a jitclass:

    import numpy as np
    from numba import int32, float32    # import the types
    from numba.experimental import jitclass

    spec = [
        ('value', int32),               # a simple scalar field
        ('array', float32[:]),          # an array field
    ]

    @jitclass(spec)
    class Bag(object):
        def __init__(self, value):
            self.value = value
            self.array = np.zeros(value, dtype=np.float32)

        @property
        def size(self):
            return self.array.size

        def increment(self, val):
            for i in range(self.size):
                self.array[i] += val
            return self.array

        @staticmethod
        def add(x, y):
            return x + y

    n = 21
    mybag = Bag(n)

In the above example, a `spec` is provided as a list of 2-tuples. The tuples contain the name of the field and the Numba type of the field. Alternatively, user can use a dictionary (an `OrderedDict` preferably for stable field ordering), which maps field names to types.

The definition of the class requires at least a `__init__` method for initializing each defined fields. Uninitialized fields contains garbage data. Methods and properties (getters and setters only) can be defined. They will be automatically compiled.

## Inferred class member types from type annotations with `as_numba_type`

Fields of a `jitclass` can also be inferred from Python type annotations.

    from typing import List
    from numba.experimental import jitclass
    from numba.typed import List as NumbaList

    @jitclass
    class Counter:
        value: int

        def __init__(self):
            self.value = 0

        def get(self) -> int:
            ret = self.value
            self.value += 1
            return ret

    @jitclass
    class ListLoopIterator:
        counter: Counter
        items: List[float]

        def __init__(self, items: List[float]):
            self.items = items
            self.counter = Counter()

        def get(self) -> float:
            idx = self.counter.get() % len(self.items)
            return self.items[idx]

    items = NumbaList([3.14, 2.718, 0.123, -4.])
    loop_itr = ListLoopIterator(items)

Any type annotations on the class will be used to extend the spec if that field is not already present. The Numba type corresponding to the given Python type is inferred using `as_numba_type`. For example, if we have the class

    @jitclass([("w", int32), ("y", float64[:])])
    class Foo:
        w: int
        x: float
        y: np.ndarray
        z: SomeOtherType

        def __init__(self, w: int, x: float, y: np.ndarray, z: SomeOtherType):
            ...

then the full spec used for `Foo` will be:

- `"w": int32` (specified in the `spec`)

- `"x": float64` (added from type annotation)

- `"y": array(float64, 1d, A)` (specified in the `spec`)

- `"z": numba.as_numba_type(SomeOtherType)` (added from type annotation)

Here `SomeOtherType` could be any supported Python type (e.g. `bool`, `typing.Dict[int, typing.Tuple[float, float]]`, or another `jitclass`).

Note that only type annotations on the class will be used to infer spec elements. Method type annotations (e.g. those of `__init__` above) are ignored.

Numba requires knowing the dtype and rank of NumPy arrays, which cannot currently be expressed with type annotations. Because of this, NumPy arrays need to be included in the `spec` explicitly.

## Specifying `numba.typed` containers as class members explicitly

The following patterns demonstrate how to specify a `numba.typed.Dict` or `numba.typed.List` explicitly as part of the `spec` passed to `jitclass`.

First, using explicit Numba types and explicit construction.

    from numba import types, typed
    from numba.experimental import jitclass

    # key and value types
    kv_ty = (types.int64, types.unicode_type)

    # A container class with:
    # * member 'd' holding a typed dictionary of int64 -> unicode string (kv_ty)
    # * member 'l' holding a typed list of float64
    @jitclass([('d', types.DictType(*kv_ty)),
               ('l', types.ListType(types.float64))])
    class ContainerHolder(object):
        def __init__(self):
            # initialize the containers
            self.d = typed.Dict.empty(*kv_ty)
            self.l = typed.List.empty_list(types.float64)

    container = ContainerHolder()
    container.d[1] = "apple"
    container.d[2] = "orange"
    container.l.append(123.)
    container.l.append(456.)
    print(container.d) # {1: apple, 2: orange}
    print(container.l) # [123.0, 456.0]

Another useful pattern is to use the `numba.typed` container attribute `_numba_type_` to find the type of a container, this can be accessed directly from an instance of the container in the Python interpreter. The same information can be obtained by calling `numba.typeof()` on the instance. For example:

    from numba import typed, typeof
    from numba.experimental import jitclass

    d = typed.Dict()
    d[1] = "apple"
    d[2] = "orange"
    l = typed.List()
    l.append(123.)
    l.append(456.)

    @jitclass([('d', typeof(d)), ('l', typeof(l))])
    class ContainerInstHolder(object):
        def __init__(self, dict_inst, list_inst):
            self.d = dict_inst
            self.l = list_inst

    container = ContainerInstHolder(d, l)
    print(container.d) # {1: apple, 2: orange}
    print(container.l) # [123.0, 456.0]

It is worth noting that the instance of the container in a `jitclass` must be initialized before use, for example, this will cause an invalid memory access as `self.d` is written to without `d` being initialized as a `type.Dict` instance of the type specified.

    from numba import types
    from numba.experimental import jitclass

    dict_ty = types.DictType(types.int64, types.unicode_type)

    @jitclass([('d', dict_ty)])
    class NotInitialisingContainer(object):
        def __init__(self):
            self.d[10] = "apple" # this is invalid, `d` is not initialized

    NotInitialisingContainer() # segmentation fault/memory access violation

## Support operations

The following operations of jitclasses work in both the interpreter and Numba compiled functions:

- calling the jitclass class object to construct a new instance (e.g. `mybag = Bag(123)`);

- read/write access to attributes and properties (e.g. `mybag.value`);

- calling methods (e.g. `mybag.increment(3)`);

- calling static methods as instance attributes (e.g. `mybag.add(1, 1)`);

- calling static methods as class attributes (e.g. `Bag.add(1, 2)`);

- using select dunder methods (e.g. `__add__` with `mybag + otherbag`);

Using jitclasses in Numba compiled function is more efficient. Short methods can be inlined (at the discretion of LLVM inliner). Attributes access are simply reading from a C structure. Using jitclasses from the interpreter has the same overhead of calling any Numba compiled function from the interpreter. Arguments and return values must be unboxed or boxed between Python objects and native representation. Values encapsulated by a jitclass does not get boxed into Python object when the jitclass instance is handed to the interpreter. It is during attribute access to the field values that they are boxed. Calling static methods as class attributes is only supported outside of the class definition (i.e. code cannot call `Bag.add()` from within another method of `Bag`).

### Supported dunder methods

The following dunder methods may be defined for jitclasses:

- `__abs__`

- `__bool__`

- `__complex__`

- `__contains__`

- `__float__`

- `__getitem__`

- `__hash__`

- `__index__`

- `__invert__`

- `__int__`

- `__len__`

- `__setitem__`

- `__str__`

- `__eq__`

- `__ne__`

- `__ge__`

- `__gt__`

- `__le__`

- `__lt__`

- `__add__`

- `__floordiv__`

- `__lshift__`

- `__matmul__`

- `__mod__`

- `__mul__`

- `__neg__`

- `__pos__`

- `__pow__`

- `__rshift__`

- `__sub__`

- `__truediv__`

- `__and__`

- `__or__`

- `__xor__`

- `__iadd__`

- `__ifloordiv__`

- `__ilshift__`

- `__imatmul__`

- `__imod__`

- `__imul__`

- `__ipow__`

- `__irshift__`

- `__isub__`

- `__itruediv__`

- `__iand__`

- `__ior__`

- `__ixor__`

- `__radd__`

- `__rfloordiv__`

- `__rlshift__`

- `__rmatmul__`

- `__rmod__`

- `__rmul__`

- `__rpow__`

- `__rrshift__`

- `__rsub__`

- `__rtruediv__`

- `__rand__`

- `__ror__`

- `__rxor__`

Refer to the Python Data Model documentation for descriptions of these methods.

## Limitations

- A jitclass class object is treated as a function (the constructor) inside a Numba compiled function.

- `isinstance()` only works in the interpreter.

- Manipulating jitclass instances in the interpreter is not optimized, yet.

- Support for jitclasses are available on CPU only. (Note: Support for GPU devices is planned for a future release.)

## The decorator: `@jitclass`

numba.experimental.jitclass(cls_or_spec=None, spec=None)

  A function for creating a jitclass. Can be used as a decorator or function.

  Different use cases will cause different arguments to be set.

  If specified, `spec` gives the types of class fields. It must be a dictionary or sequence. With a dictionary, use collections.OrderedDict for stable ordering. With a sequence, it must contain 2-tuples of (fieldname, fieldtype).

  Any class annotations for field names not listed in spec will be added. For class annotation x: T we will append `("x", as_numba_type(T))` to the spec if `x` is not already a key in spec.

  Returns:
    If used as a decorator, returns a callable that takes a class object and
    returns a compiled version.
    If used as a function, returns the compiled class (an instance of
    `JitClassType`).

  Examples

  1. `cls_or_spec = None`, `spec = None`

      >>> @jitclass()
      ... class Foo:
      ...     ...

  2. `cls_or_spec = None`, `spec = spec`

      >>> @jitclass(spec=spec)
      ... class Foo:
      ...     ...

  3. `cls_or_spec = Foo`, `spec = None`

      >>> @jitclass
      ... class Foo:
      ...     ...

  4) `cls_or_spec = spec`, `spec = None` In this case we update `cls_or_spec, spec = None, cls_or_spec`.

      >>> @jitclass(spec)
      ... class Foo:
      ...     ...

  5. `cls_or_spec = Foo`, `spec = spec`

      >>> JitFoo = jitclass(Foo, spec)