2111476234

Build Type

github

Committed by GitHub

Commit Message

Merge 6eeabfaea into d9e692141

Pull Request Pull Request #177: WIP: begin adding user-defined types

Run Details

4630 of 4671 branches covered (99.12%)

Branch coverage included in aggregate %.

861 of 908 new or added lines in 16 files covered. (94.82%)

1 existing line in 1 file now uncovered.

14623 of 14679 relevant lines covered (99.62%)

14.57 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

94.22

/grblas/dtypes.py

import numba
import numpy as np
from numpy import find_common_type, promote_types

from . import ffi, lib

# Default assumption unless FC32/FC64 are found in lib
_supports_complex = hasattr(lib, "GrB_FC64") or hasattr(lib, "GxB_FC64")


class DataType:
    __slots__ = "name", "gb_obj", "gb_name", "c_type", "numba_type", "np_type", "__weakref__"

    def __init__(self, name, gb_obj, gb_name, c_type, numba_type, np_type):
        self.name = name
        self.gb_obj = gb_obj
        self.gb_name = gb_name
        self.c_type = c_type
        self.numba_type = numba_type
        self.np_type = np.dtype(np_type)

    def __repr__(self):
        return self.name

    def __eq__(self, other):
        if type(other) is DataType:
            return self is other
        # Attempt to use `other` as a lookup key
        try:
            other = lookup_dtype(other)
            return self is other
        except ValueError:
            raise TypeError(f"Invalid or unknown datatype: {other}") from None

    def __hash__(self):
        return hash(self.gb_obj)

    def __lt__(self, other):
        # Let us sort for prettier error reporting
        if type(other) is DataType:
            t1 = self.np_type
            t2 = other.np_type
            return (t1.kind, t1.itemsize, t1.name) < (t2.kind, t2.itemsize, t2.name)
        # Attempt to use `other` as a lookup key
        try:
            return self < lookup_dtype(other)
        except ValueError:
            raise TypeError(f"Invalid or unknown datatype: {other}") from None

    def __reduce__(self):
        if self._is_udt:
            1 / 0
            return (self._deserialize, (self.name, self.np_type, self._is_anonymous))
        if self.gb_name == "GrB_Index":
            return "_INDEX"
        return self.name

    @property
    def _carg(self):
        return self.gb_obj[0] if self.gb_name is None else self.gb_obj

    @property
    def _is_anonymous(self):
        return globals().get(self.name) is not self

    @property
    def _is_udt(self):
        return self.gb_name is None

    @staticmethod
    def _deserialize(name, dtype, is_anonymous):
        1 / 0
        if is_anonymous:
            return register_anonymous(dtype, name)
        if name in _registry:
            return _registry[name]
        return register_new(name, dtype)


def register_new(name, dtype):
    if not name.isidentifier():
        raise ValueError(f"`name` argument must be a valid Python identifier; got: {name!r}")
    if name in _registry or name in globals():
        raise ValueError(f"{name!r} name for dtype is unavailable")
    rv = register_anonymous(dtype, name)
    _registry[name] = rv
    globals()[name] = rv
    return rv


def register_anonymous(dtype, name=None):
    from .exceptions import check_status_carg

    dtype = np.dtype(dtype)
    if dtype.hasobject:
        raise ValueError("dtype must not allow Python objects")
    if dtype.isbuiltin != 0:
        raise ValueError("dtype must not be a builtin type")
    if name is None:
        name = repr(dtype)
    numba_type = numba.typeof(dtype).dtype

    if dtype in _registry:
        gb_obj = _registry[dtype].gb_obj
    else:
        gb_obj = ffi.new("GrB_Type*")
        status = lib.GrB_Type_new(gb_obj, dtype.itemsize)
        check_status_carg(status, "Type", gb_obj[0])
    # For now, let's use "opaque" unsigned bytes for the c type.
    # grb_name probably isn't useful, right?
    rv = DataType(name, gb_obj, None, f"uint8_t[{dtype.itemsize}]", numba_type, dtype)
    if dtype not in _registry:
        _registry[gb_obj] = rv
        _registry[dtype] = rv
        _registry[numba_type] = rv
        _registry[numba_type.name] = rv
    return rv


BOOL = DataType("BOOL", lib.GrB_BOOL, "GrB_BOOL", "_Bool", numba.types.bool_, np.bool_)
INT8 = DataType("INT8", lib.GrB_INT8, "GrB_INT8", "int8_t", numba.types.int8, np.int8)
UINT8 = DataType("UINT8", lib.GrB_UINT8, "GrB_UINT8", "uint8_t", numba.types.uint8, np.uint8)
INT16 = DataType("INT16", lib.GrB_INT16, "GrB_INT16", "int16_t", numba.types.int16, np.int16)
UINT16 = DataType("UINT16", lib.GrB_UINT16, "GrB_UINT16", "uint16_t", numba.types.uint16, np.uint16)
INT32 = DataType("INT32", lib.GrB_INT32, "GrB_INT32", "int32_t", numba.types.int32, np.int32)
UINT32 = DataType("UINT32", lib.GrB_UINT32, "GrB_UINT32", "uint32_t", numba.types.uint32, np.uint32)
INT64 = DataType("INT64", lib.GrB_INT64, "GrB_INT64", "int64_t", numba.types.int64, np.int64)
# _Index (like UINT64) is for internal use only and shouldn't be exposed to the user
_INDEX = DataType("UINT64", lib.GrB_UINT64, "GrB_Index", "GrB_Index", numba.types.uint64, np.uint64)
UINT64 = DataType("UINT64", lib.GrB_UINT64, "GrB_UINT64", "uint64_t", numba.types.uint64, np.uint64)
FP32 = DataType("FP32", lib.GrB_FP32, "GrB_FP32", "float", numba.types.float32, np.float32)
FP64 = DataType("FP64", lib.GrB_FP64, "GrB_FP64", "double", numba.types.float64, np.float64)

if _supports_complex and hasattr(lib, "GxB_FC32"):
    FC32 = DataType(
        "FC32", lib.GxB_FC32, "GxB_FC32", "float _Complex", numba.types.complex64, np.complex64
    )
if _supports_complex and hasattr(lib, "GrB_FC32"):  # pragma: no coverage
    FC32 = DataType(
        "FC32", lib.GrB_FC32, "GrB_FC32", "float _Complex", numba.types.complex64, np.complex64
    )
if _supports_complex and hasattr(lib, "GxB_FC64"):
    FC64 = DataType(
        "FC64", lib.GxB_FC64, "GxB_FC64", "double _Complex", numba.types.complex128, np.complex128
    )
if _supports_complex and hasattr(lib, "GrB_FC64"):  # pragma: no coverage
    FC64 = DataType(
        "FC64", lib.GrB_FC64, "GrB_FC64", "double _Complex", numba.types.complex128, np.complex128
    )

# Used for testing user-defined functions
_sample_values = {
    INT8: np.int8(1),
    UINT8: np.uint8(1),
    INT16: np.int16(1),
    UINT16: np.uint16(1),
    INT32: np.int32(1),
    UINT32: np.uint32(1),
    INT64: np.int64(1),
    UINT64: np.uint64(1),
    FP32: np.float32(0.5),
    FP64: np.float64(0.5),
    BOOL: np.bool_(True),
}
if _supports_complex:
    _sample_values.update(
        {
            FC32: np.complex64(complex(0, 0.5)),
            FC64: np.complex128(complex(0, 0.5)),
        }
    )

# Create register to easily lookup types by name, gb_obj, or c_type
_registry = {}
_dtypes_to_register = [
    BOOL,
    INT8,
    UINT8,
    INT16,
    UINT16,
    INT32,
    UINT32,
    INT64,
    UINT64,
    FP32,
    FP64,
]
if _supports_complex:
    _dtypes_to_register.extend([FC32, FC64])

for dtype in _dtypes_to_register:
    _registry[dtype.name] = dtype
    _registry[dtype.name.lower()] = dtype
    _registry[dtype.gb_obj] = dtype
    _registry[dtype.gb_name] = dtype
    _registry[dtype.gb_name.lower()] = dtype
    _registry[dtype.c_type] = dtype
    _registry[dtype.c_type.upper()] = dtype
    _registry[dtype.numba_type] = dtype
    _registry[dtype.numba_type.name] = dtype
    val = _sample_values[dtype]
    _registry[val.dtype] = dtype
    _registry[val.dtype.name] = dtype
# Upcast numpy float16 to float32
_registry[np.dtype(np.float16)] = FP32
_registry["float16"] = FP32

# Add some common Python types as lookup keys
_registry[bool] = BOOL
_registry[int] = INT64
_registry[float] = FP64
_registry["bool"] = BOOL
_registry["int"] = INT64
_registry["float"] = FP64  # Choose 'float' to match numpy/Python; c_type 'float' would be FP32
if _supports_complex:
    _registry[complex] = FC64
    _registry["complex"] = FC64


def lookup_dtype(key, value=None):
    # Check for silly lookup where key is already a DataType
    if type(key) is DataType:
        return key
    try:
        return _registry[key]
    except KeyError:
        pass
    if hasattr(key, "name"):
        try:
            return _registry[key.name]
        except KeyError:
            pass
    if value is not None and hasattr(value, "dtype") and value.dtype in _registry:
        return _registry[value.dtype]
    try:
        return lookup_dtype(np.dtype(key))
    except Exception:
        pass
    try:
        return lookup_dtype(key.literal_type)  # For numba dtype inference
    except Exception:
        pass
    raise ValueError(f"Unknown dtype: {key}", type(key))


def unify(type1, type2, *, is_left_scalar=False, is_right_scalar=False):
    """
    Returns a type that can hold both type1 and type2

    For example:
    unify(INT32, INT64) -> INT64
    unify(INT8, UINT16) -> INT32
    unify(BOOL, UINT16) -> UINT16
    unify(FP32, INT32) -> FP64
    """
    if type1 is type2:
        return type1
    if is_left_scalar:
        scalar_types = [type1.np_type]
        array_types = []
    elif not is_right_scalar:
        # Using `promote_types` is faster than `find_common_type`
        return lookup_dtype(promote_types(type1.np_type, type2.np_type))
    else:
        scalar_types = []
        array_types = [type1.np_type]
    if is_right_scalar:
        scalar_types.append(type2.np_type)
    else:
        array_types.append(type2.np_type)
    return lookup_dtype(find_common_type(array_types, scalar_types))

1	import numba	15✔
2	import numpy as np	15✔
3	from numpy import find_common_type, promote_types	15✔
4
5	from . import ffi, lib	15✔
6
7	# Default assumption unless FC32/FC64 are found in lib
8	_supports_complex = hasattr(lib, "GrB_FC64") or hasattr(lib, "GxB_FC64")	15✔
9
10
11	class DataType:	15✔
12	__slots__ = "name", "gb_obj", "gb_name", "c_type", "numba_type", "np_type", "__weakref__"	15✔
13
14	def __init__(self, name, gb_obj, gb_name, c_type, numba_type, np_type):	15✔
15	self.name = name	15✔
16	self.gb_obj = gb_obj	15✔
17	self.gb_name = gb_name	15✔
18	self.c_type = c_type	15✔
19	self.numba_type = numba_type	15✔
20	self.np_type = np.dtype(np_type)	15✔
21
22	def __repr__(self):	15✔
23	return self.name	15✔
24
25	def __eq__(self, other):	15✔
26	if type(other) is DataType:	15✔
27	return self is other	15✔
28	# Attempt to use `other` as a lookup key
29	try:	15✔
30	other = lookup_dtype(other)	15✔
31	return self is other	15✔
32	except ValueError:	15✔
33	raise TypeError(f"Invalid or unknown datatype: {other}") from None	15✔
34
35	def __hash__(self):	15✔
36	return hash(self.gb_obj)	15✔
37
38	def __lt__(self, other):	15✔
39	# Let us sort for prettier error reporting
40	if type(other) is DataType:	15✔
41	t1 = self.np_type	15✔
42	t2 = other.np_type	15✔
43	return (t1.kind, t1.itemsize, t1.name) < (t2.kind, t2.itemsize, t2.name)	15✔
44	# Attempt to use `other` as a lookup key
45	try:	15✔
46	return self < lookup_dtype(other)	15✔
47	except ValueError:	15✔
48	raise TypeError(f"Invalid or unknown datatype: {other}") from None	15✔
49
50	def __reduce__(self):	15✔
51	if self._is_udt:	15!
NEW 52	1 / 0	×
NEW 53	return (self._deserialize, (self.name, self.np_type, self._is_anonymous))	×
54	if self.gb_name == "GrB_Index":	15✔
55	return "_INDEX"	15✔
56	return self.name	15✔
57
58	@property	15✔
59	def _carg(self):	15✔
60	return self.gb_obj[0] if self.gb_name is None else self.gb_obj	15✔
61
62	@property	15✔
63	def _is_anonymous(self):	15✔
64	return globals().get(self.name) is not self	15✔
65
66	@property	15✔
67	def _is_udt(self):	15✔
68	return self.gb_name is None	15✔
69
70	@staticmethod	15✔
71	def _deserialize(name, dtype, is_anonymous):	15✔
NEW 72	1 / 0	×
NEW 73	if is_anonymous:	×
NEW 74	return register_anonymous(dtype, name)	×
NEW 75	if name in _registry:	×
NEW 76	return _registry[name]	×
NEW 77	return register_new(name, dtype)	×
78
79
80	def register_new(name, dtype):	15✔
81	if not name.isidentifier():	15✔
82	raise ValueError(f"`name` argument must be a valid Python identifier; got: {name!r}")	15✔
83	if name in _registry or name in globals():	15✔
84	raise ValueError(f"{name!r} name for dtype is unavailable")	15✔
85	rv = register_anonymous(dtype, name)	15✔
86	_registry[name] = rv	15✔
87	globals()[name] = rv	15✔
88	return rv	15✔
89
90
91	def register_anonymous(dtype, name=None):	15✔
92	from .exceptions import check_status_carg	15✔
93
94	dtype = np.dtype(dtype)	15✔
95	if dtype.hasobject:	15✔
96	raise ValueError("dtype must not allow Python objects")	15✔
97	if dtype.isbuiltin != 0:	15✔
98	raise ValueError("dtype must not be a builtin type")	15✔
99	if name is None:	15✔
100	name = repr(dtype)	15✔
101	numba_type = numba.typeof(dtype).dtype	15✔
102
103	if dtype in _registry:	15✔
104	gb_obj = _registry[dtype].gb_obj	15✔
105	else:
106	gb_obj = ffi.new("GrB_Type*")	15✔
107	status = lib.GrB_Type_new(gb_obj, dtype.itemsize)	15✔
108	check_status_carg(status, "Type", gb_obj[0])	15✔
109	# For now, let's use "opaque" unsigned bytes for the c type.
110	# grb_name probably isn't useful, right?
111	rv = DataType(name, gb_obj, None, f"uint8_t[{dtype.itemsize}]", numba_type, dtype)	15✔
112	if dtype not in _registry:	15✔
113	_registry[gb_obj] = rv	15✔
114	_registry[dtype] = rv	15✔
115	_registry[numba_type] = rv	15✔
116	_registry[numba_type.name] = rv	15✔
117	return rv	15✔
118
119
120	BOOL = DataType("BOOL", lib.GrB_BOOL, "GrB_BOOL", "_Bool", numba.types.bool_, np.bool_)	15✔
121	INT8 = DataType("INT8", lib.GrB_INT8, "GrB_INT8", "int8_t", numba.types.int8, np.int8)	15✔
122	UINT8 = DataType("UINT8", lib.GrB_UINT8, "GrB_UINT8", "uint8_t", numba.types.uint8, np.uint8)	15✔
123	INT16 = DataType("INT16", lib.GrB_INT16, "GrB_INT16", "int16_t", numba.types.int16, np.int16)	15✔
124	UINT16 = DataType("UINT16", lib.GrB_UINT16, "GrB_UINT16", "uint16_t", numba.types.uint16, np.uint16)	15✔
125	INT32 = DataType("INT32", lib.GrB_INT32, "GrB_INT32", "int32_t", numba.types.int32, np.int32)	15✔
126	UINT32 = DataType("UINT32", lib.GrB_UINT32, "GrB_UINT32", "uint32_t", numba.types.uint32, np.uint32)	15✔
127	INT64 = DataType("INT64", lib.GrB_INT64, "GrB_INT64", "int64_t", numba.types.int64, np.int64)	15✔
128	# _Index (like UINT64) is for internal use only and shouldn't be exposed to the user
129	_INDEX = DataType("UINT64", lib.GrB_UINT64, "GrB_Index", "GrB_Index", numba.types.uint64, np.uint64)	15✔
130	UINT64 = DataType("UINT64", lib.GrB_UINT64, "GrB_UINT64", "uint64_t", numba.types.uint64, np.uint64)	15✔
131	FP32 = DataType("FP32", lib.GrB_FP32, "GrB_FP32", "float", numba.types.float32, np.float32)	15✔
132	FP64 = DataType("FP64", lib.GrB_FP64, "GrB_FP64", "double", numba.types.float64, np.float64)	15✔
133
134	if _supports_complex and hasattr(lib, "GxB_FC32"):	15✔
135	FC32 = DataType(	11✔
136	"FC32", lib.GxB_FC32, "GxB_FC32", "float _Complex", numba.types.complex64, np.complex64
137	)
138	if _supports_complex and hasattr(lib, "GrB_FC32"): # pragma: no coverage
139	FC32 = DataType(
140	"FC32", lib.GrB_FC32, "GrB_FC32", "float _Complex", numba.types.complex64, np.complex64
141	)
142	if _supports_complex and hasattr(lib, "GxB_FC64"):	15✔
143	FC64 = DataType(	11✔
144	"FC64", lib.GxB_FC64, "GxB_FC64", "double _Complex", numba.types.complex128, np.complex128
145	)
146	if _supports_complex and hasattr(lib, "GrB_FC64"): # pragma: no coverage
147	FC64 = DataType(
148	"FC64", lib.GrB_FC64, "GrB_FC64", "double _Complex", numba.types.complex128, np.complex128
149	)
150
151	# Used for testing user-defined functions
152	_sample_values = {	15✔
153	INT8: np.int8(1),
154	UINT8: np.uint8(1),
155	INT16: np.int16(1),
156	UINT16: np.uint16(1),
157	INT32: np.int32(1),
158	UINT32: np.uint32(1),
159	INT64: np.int64(1),
160	UINT64: np.uint64(1),
161	FP32: np.float32(0.5),
162	FP64: np.float64(0.5),
163	BOOL: np.bool_(True),
164	}
165	if _supports_complex:	15✔
166	_sample_values.update(	11✔
167	{
168	FC32: np.complex64(complex(0, 0.5)),
169	FC64: np.complex128(complex(0, 0.5)),
170	}
171	)
172
173	# Create register to easily lookup types by name, gb_obj, or c_type
174	_registry = {}	15✔
175	_dtypes_to_register = [	15✔
176	BOOL,
177	INT8,
178	UINT8,
179	INT16,
180	UINT16,
181	INT32,
182	UINT32,
183	INT64,
184	UINT64,
185	FP32,
186	FP64,
187	]
188	if _supports_complex:	15✔
189	_dtypes_to_register.extend([FC32, FC64])	11✔
190
191	for dtype in _dtypes_to_register:	15✔
192	_registry[dtype.name] = dtype	15✔
193	_registry[dtype.name.lower()] = dtype	15✔
194	_registry[dtype.gb_obj] = dtype	15✔
195	_registry[dtype.gb_name] = dtype	15✔
196	_registry[dtype.gb_name.lower()] = dtype	15✔
197	_registry[dtype.c_type] = dtype	15✔
198	_registry[dtype.c_type.upper()] = dtype	15✔
199	_registry[dtype.numba_type] = dtype	15✔
200	_registry[dtype.numba_type.name] = dtype	15✔
201	val = _sample_values[dtype]	15✔
202	_registry[val.dtype] = dtype	15✔
203	_registry[val.dtype.name] = dtype	15✔
204	# Upcast numpy float16 to float32
205	_registry[np.dtype(np.float16)] = FP32	15✔
206	_registry["float16"] = FP32	15✔
207
208	# Add some common Python types as lookup keys
209	_registry[bool] = BOOL	15✔
210	_registry[int] = INT64	15✔
211	_registry[float] = FP64	15✔
212	_registry["bool"] = BOOL	15✔
213	_registry["int"] = INT64	15✔
214	_registry["float"] = FP64 # Choose 'float' to match numpy/Python; c_type 'float' would be FP32	15✔
215	if _supports_complex:	15✔
216	_registry[complex] = FC64	11✔
217	_registry["complex"] = FC64	11✔
218
219
220	def lookup_dtype(key, value=None):	15✔
221	# Check for silly lookup where key is already a DataType
222	if type(key) is DataType:	15✔
223	return key	15✔
224	try:	15✔
225	return _registry[key]	15✔
226	except KeyError:	15✔
227	pass	15✔
228	if hasattr(key, "name"):	15✔
229	try:	15✔
230	return _registry[key.name]	15✔
231	except KeyError:	15✔
232	pass	15✔
233	if value is not None and hasattr(value, "dtype") and value.dtype in _registry:	15✔
234	return _registry[value.dtype]	15✔
235	try:	15✔
236	return lookup_dtype(np.dtype(key))	15✔
237	except Exception:	15✔
238	pass	15✔
239	try:	15✔
240	return lookup_dtype(key.literal_type) # For numba dtype inference	15✔
241	except Exception:	15✔
242	pass	15✔
243	raise ValueError(f"Unknown dtype: {key}", type(key))	15✔
244
245
246	def unify(type1, type2, *, is_left_scalar=False, is_right_scalar=False):	15✔
247	"""
248	Returns a type that can hold both type1 and type2
249
250	For example:
251	unify(INT32, INT64) -> INT64
252	unify(INT8, UINT16) -> INT32
253	unify(BOOL, UINT16) -> UINT16
254	unify(FP32, INT32) -> FP64
255	"""
256	if type1 is type2:	15✔
257	return type1	15✔
258	if is_left_scalar:	15✔
259	scalar_types = [type1.np_type]	15✔
260	array_types = []	15✔
261	elif not is_right_scalar:	15✔
262	# Using `promote_types` is faster than `find_common_type`
263	return lookup_dtype(promote_types(type1.np_type, type2.np_type))	15✔
264	else:
265	scalar_types = []	15✔
266	array_types = [type1.np_type]	15✔
267	if is_right_scalar:	15✔
268	scalar_types.append(type2.np_type)	15✔
269	else:
270	array_types.append(type2.np_type)	15✔
271	return lookup_dtype(find_common_type(array_types, scalar_types))	15✔

metagraph-dev / grblas / 2111476234

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous