diff --git a/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_rvec.py b/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_rvec.py index d92b0f2c3ab57..0814eafe5eb47 100644 --- a/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_rvec.py +++ b/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_rvec.py @@ -62,36 +62,27 @@ from . import pythonization -_array_interface_dtype_map = { - "Long64_t": "i", - "ULong64_t": "u", - "double": "f", - "float": "f", - "int": "i", - "long": "i", - "unsigned char": "b", - "unsigned int": "u", - "unsigned long": "u", + +_numpy_dtype_typeinfo_map = { + "i1": {"typedef": "Char_t", "cpp": "char"}, + "u1": {"typedef": "UChar_t", "cpp": "unsigned char"}, + "i2": {"typedef": "Short_t", "cpp": "short"}, + "u2": {"typedef": "UShort_t", "cpp": "unsigned short"}, + "i4": {"typedef": "Int_t", "cpp": "int"}, + "u4": {"typedef": "UInt_t", "cpp": "unsigned int"}, + "i8": {"typedef": "Long64_t", "cpp": "long long"}, + "u8": {"typedef": "ULong64_t", "cpp": "unsigned long long"}, + "f4": {"typedef": "Float_t", "cpp": "float"}, + "f8": {"typedef": "Double_t", "cpp": "double"}, + "b1": {"typedef": "Bool_t", "cpp": "bool"} } def _get_cpp_type_from_numpy_type(dtype): - cpptypes = { - "i2": "Short_t", - "u2": "UShort_t", - "i4": "int", - "u4": "unsigned int", - "i8": "Long64_t", - "u8": "ULong64_t", - "f4": "float", - "f8": "double", - "b1": "bool", - } - - if dtype not in cpptypes: + if dtype not in _numpy_dtype_typeinfo_map: raise RuntimeError("Object not convertible: Python object has unknown data-type '" + dtype + "'.") - return cpptypes[dtype] + return _numpy_dtype_typeinfo_map[dtype]['cpp'] def _AsRVec(arr): @@ -160,10 +151,11 @@ def _AsRVec(arr): def get_array_interface(self): cppname = type(self).__cpp_name__ - for dtype in _array_interface_dtype_map: - if cppname.endswith("<{}>".format(dtype)): - dtype_numpy = _array_interface_dtype_map[dtype] - dtype_size = cppyy.sizeof(dtype) + for np_dtype, type_info in _numpy_dtype_typeinfo_map.items(): + root_typedef = type_info["typedef"] + cpp_type = type_info["cpp"] + if cppname.endswith("<{}>".format(root_typedef)) or cppname.endswith("<{}>".format(cpp_type)): + dtype_numpy = np_dtype endianness = "<" if sys.byteorder == "little" else ">" size = self.size() # Numpy breaks for data pointer of 0 even though the array is empty. @@ -174,14 +166,17 @@ def get_array_interface(self): pointer = cppyy.ll.addressof(self.data()) return { "shape": (size,), - "typestr": "{}{}{}".format(endianness, dtype_numpy, dtype_size), + "typestr": "{}{}".format(endianness, dtype_numpy), "version": 3, "data": (pointer, False), } def add_array_interface_property(klass, name): - if True in [name.endswith("<{}>".format(dtype)) for dtype in _array_interface_dtype_map]: + if any( + name.endswith(f"<{type_info['typedef']}>") or name.endswith(f"<{type_info['cpp']}>") + for type_info in _numpy_dtype_typeinfo_map.values() + ): klass.__array_interface__ = property(get_array_interface) diff --git a/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_tmva/_rtensor.py b/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_tmva/_rtensor.py index 7fad933c93723..6a89b97e6691e 100644 --- a/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_tmva/_rtensor.py +++ b/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_tmva/_rtensor.py @@ -8,7 +8,7 @@ # For the list of contributors see $ROOTSYS/README/CREDITS. # ################################################################################ from .. import pythonization -from .._rvec import _array_interface_dtype_map, _get_cpp_type_from_numpy_type +from .._rvec import _numpy_dtype_typeinfo_map, _get_cpp_type_from_numpy_type import cppyy import sys @@ -85,8 +85,15 @@ def get_array_interface(self): idx1 = cppname.find("RTensor<") idx2 = cppname.find(",", idx1) dtype = cppname[idx1 + 8 : idx2] - dtype_numpy = _array_interface_dtype_map[dtype] - dtype_size = cppyy.sizeof(dtype) + for numpy_dtype, type_info in _numpy_dtype_typeinfo_map.items(): + if dtype == type_info["cpp"] or dtype == type_info["typedef"]: + dtype_numpy = numpy_dtype + dtype_size = cppyy.sizeof(dtype) + break + + if dtype_numpy: + raise RuntimeError(f"Unsupported dtype '{dtype}' found in RTensor.") + endianness = "<" if sys.byteorder == "little" else ">" shape = self.GetShape() strides = self.GetStrides() @@ -98,7 +105,7 @@ def get_array_interface(self): return { "shape": tuple(s for s in shape), "strides": tuple(s * dtype_size for s in strides), - "typestr": "{}{}{}".format(endianness, dtype_numpy, dtype_size), + "typestr": "{}{}".format(endianness, dtype_numpy), "version": 3, "data": (pointer, False), } @@ -113,7 +120,10 @@ def add_array_interface_property(klass, name): klass: class to be pythonized name: string containing the name of the class """ - if True in [name.find("RTensor<{},".format(dtype)) != -1 for dtype in _array_interface_dtype_map]: + if any( + ("RTensor<{},".format(type_info["cpp"]) in name) or ("RTensor<{},".format(type_info["typedef"]) in name) + for type_info in _numpy_dtype_typeinfo_map.values() + ): klass.__array_interface__ = property(get_array_interface) diff --git a/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_ttree.py b/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_ttree.py index 23dbecdc8b15a..b93f2dc53fc4b 100644 --- a/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_ttree.py +++ b/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_ttree.py @@ -160,7 +160,7 @@ """ from ROOT.libROOTPythonizations import GetBranchAttr, BranchPyz -from ._rvec import _array_interface_dtype_map, _get_cpp_type_from_numpy_type +from ._rvec import _get_cpp_type_from_numpy_type from . import pythonization from ROOT._pythonization._memory_utils import _should_give_up_ownership, _constructor_releasing_ownership, _SetDirectory_SetOwnership