def test_missing_source(self):
@structref.register
class ParticleType(types.StructRef):
pass
class Particle(structref.StructRefProxy):
def __new__(cls, pos, mass):
return structref.StructRefProxy.__new__(cls, pos)
# didn't provide the required mass argument ----^
structref.define_proxy(Particle, ParticleType, ["pos", "mass"])
with self.assertRaises(errors.TypingError) as raises:
Particle(pos=1, mass=2)
excstr = str(raises.exception)
self.assertIn("required positional argument: 'mass'", excstr)
return self.flags
@numba.njit
def Heap_get_distances(self):
return self.distances
@numba.njit
def Heap_get_indices(self):
return self.indices
structref.define_proxy(
Heap,
HeapType,
["indices", "distances", "flags"],
)
# Heap = namedtuple("Heap", ("indices", "distances", "flags"))
@numba.njit()
def make_heap(n_points, size):
"""Constructor for the numba enabled heap objects. The heaps are used
for approximate nearest neighbor search, maintaining a list of potential
neighbors sorted by their distance. We also flag if potential neighbors
are newly added to the list or not. Internally this is stored as
a single ndarray; the first axis determines whether we are looking at the
array of candidate graph_indices, the array of distances, or the flag array for
whether elements are new or not. Each of these arrays are of shape
@structref.register
class MyStructType(types.StructRef):
"""Test associated with this type represent the higher-level uses of
structef.
"""
pass
# Call to define_proxy is needed to register the use of `MyStruct` as a
# PyObject proxy for creating a Numba-allocated structref.
# The `MyStruct` class can then be used in both jit-code and interpreted-code.
structref.define_proxy(
MyStruct,
MyStructType,
['values', 'counter'],
)
@njit
def my_struct(values, counter):
st = structref.new(my_struct_ty)
my_struct_init(st, values, counter)
return st
@njit
def my_struct_init(self, values, counter):
self.values = values
self.counter = counter
"""
Wire together the Numba and Python types.
"""
import numpy as np
from numba.extending import overload_method
from numba.experimental import structref
from ._struct import CSRType
from .csr import CSR
from . import _rows, structure
structref.define_proxy(
CSR, CSRType, ['nrows', 'ncols', 'nnz', 'rowptrs', 'colinds', 'values'])
@overload_method(CSRType, 'row_extent')
def _csr_row_extent(csr, row):
return _rows.extent
@overload_method(CSRType, 'row')
def _csr_row(csr, row):
if csr.has_values:
return _rows._array_vals
else:
return _rows._array_ones
@overload_method(CSRType, 'row_cs')
def _csr_row_cs(csr, row):
@njit
def MyStruct_get_name(self):
# In jit-code, the StructRef's attribute is exposed via
# structref.register
return self.name
@njit
def MyStruct_get_vector(self):
return self.vector
# This associates the proxy with MyStructType for the given set of
# fields. Notice how we are not contraining the type of each field.
# Field types remain generic.
structref.define_proxy(MyStruct, MyStructType, ["name", "vector"])
# magictoken.ex_structref_type_definition.end
@skip_unless_scipy
class TestStructRefUsage(unittest.TestCase):
def test_type_definition(self):
np.random.seed(0)
# Redirect print
buf = []
def print(*args):
buf.append(args)
# magictoken.ex_structref_type_definition_test.begin
# Let's test our new StructRef.
class MyStruct(structref.StructRefProxy):
def __new__(cls, d, v):
return structref.StructRefProxy.__new__(cls, d, v)
@property
def A(self):
return MyStruct_get_A(self)
@property
def B(self):
return MyStruct_get_B(self)
structref.define_proxy(MyStruct, MyStructTypeTemplate, ['A', 'B'])
MyStructType = MyStructTypeTemplate(fields=data_fields)
#Struct Definition
print(MyStruct._numba_box_)
class UntypedStructRef(Type):
"""
Pointer to a Numba "meminfo" (i.e. the information for a managed
piece of memory).
"""
mutable = True
def __init__(self, dtype, meminfo):
@structref.register
class MyStructType(types.StructRef):
"""Test associated with this type represent the higher-level uses of
structef.
"""
pass
# Call to define_proxy is needed to register the use of `MyStruct` as a
# PyObject proxy for creating a Numba-allocated structref.
# The `MyStruct` class can then be used in both jit-code and interpreted-code.
structref.define_proxy(
MyStruct,
MyStructType,
['values', 'counter'],
)
@njit
def my_struct(values, counter):
st = structref.new(my_struct_ty)
my_struct_init(st, values, counter)
return st
@njit
def my_struct_init(self, values, counter):
self.values = values
self.counter = counter
if config.DISABLE_JIT:
@dataclasses.dataclass
class mkl_h:
H: int
nrows: int
ncols: int
csr_ref: Optional[np.ndarray]
else:
class mkl_h(structref.StructRefProxy):
"""
Type for MKL handles. Opaque, do not use directly.
"""
structref.define_proxy(mkl_h, mkl_h_type,
['H', 'nrows', 'ncols', 'csr_ref'])
def _make_handle_impl(csr):
"Make a handle from a known-constructable CSR"
_sp = ffi.from_buffer(csr.rowptrs)
_cols = ffi.from_buffer(csr.colinds)
vs = csr.values
assert vs.size == csr.nnz
_vals = ffi.from_buffer(vs)
h = lk_mkl_spcreate(csr.nrows, csr.ncols, _sp, _cols, _vals)
lk_mkl_spopt(h)
return mkl_h(h, csr.nrows, csr.ncols, csr)
_make_handle = njit(_make_handle_impl)
def modify(self, fact, attr, val):
return modify(self, fact, attr, val)
def __del__(self):
pass
# kb_data_dtor(self.kb_data)
@structref.register
class KnowledgeBaseTypeTemplate(types.StructRef):
def preprocess_fields(self, fields):
return tuple((name, types.unliteral(typ)) for name, typ in fields)
structref.define_proxy(KnowledgeBase, KnowledgeBaseTypeTemplate,
["kb_data", "context_data"])
KnowledgeBaseType = KnowledgeBaseTypeTemplate(fields=[(
"kb_data",
KnowledgeBaseDataType), ("context_data", KnowledgeBaseContextDataType)])
@njit(cache=True)
def facts_for_t_id(kb_data, t_id):
L = len(kb_data.facts)
if (t_id >= L):
expand_kb_data_types(kb_data, 1 + L - t_id)
return _struct_from_pointer(VectorType, kb_data.facts[t_id])
@njit(cache=True)
def fact_at_f_id(typ, t_id_facts, f_id):
