[WIP][Python] Unify maps used for numpy to cpp conversion

bindings/pyroot/pythonizations/python/ROOT/_pythonization/_rvec.py

@@ -56,42 +56,33 @@
 \endpythondoc
 """
 
 import sys
 
 import cppyy
 
 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 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 @@
                 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)
 
 

bindings/pyroot/pythonizations/python/ROOT/_pythonization/_tmva/_rtensor.py

@@ -7,10 +7,10 @@
 # For the licensing terms see $ROOTSYS/LICENSE.                                #
 # 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
 
 
 def _AsRTensor(arr):
@@ -24,9 +24,9 @@
     PyObject. The RTensor takes the data pointer and the shape from the array
     interface dictionary.
     """
     import ROOT
     import math
     import platform
 
     # Get array interface of object
     interface = arr.__array_interface__
@@ -36,7 +36,7 @@
 
     # Get the size of the contiguous memory
     shape = interface["shape"]
     size = math.prod(shape) if len(shape) > 0 else 0
 
     # Get the typestring and properties thereof
     typestr = interface["typestr"]
@@ -85,8 +85,15 @@
     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 @@
     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 @@
         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)
 
 
@@ -140,13 +150,13 @@
     # Convert negative indices and Nones
     isSlice = False
     for i, x in enumerate(idx):
         if type(x) == slice:
             isSlice = True
             start = 0 if x.start is None else x.start
             stop = shape[i] if x.stop is None else x.stop
             if stop < 0:
                 stop += shape[i]
             if x.step != None:
                 raise Exception("RTensor does not support slices with step size unequal 1.")
             idx[i] = slice(start, stop, None)
         else:
@@ -158,7 +168,7 @@
         idxVec = cppyy.gbl.std.vector("vector<size_t>")(len(idx))
         for i, x in enumerate(idx):
             idxVec[i].resize(2)
             if type(x) == slice:
                 idxVec[i][0] = x.start
                 idxVec[i][1] = x.stop
             else:

bindings/pyroot/pythonizations/python/ROOT/_pythonization/_ttree.py

@@ -159,10 +159,10 @@
 \endpythondoc
 """
 
 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
 
 # TTree iterator
 def _TTree__iter__(self):
@@ -181,8 +181,8 @@
     """Helper for the SetBranchAddress pythonization, extracting the relevant
     address from a Python object if possible.
     """
     import cppyy
     import ctypes
 
     is_leaf_list = branch.IsA() is cppyy.gbl.TBranch.Class()