diff --git a/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_stl_vector.py b/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_stl_vector.py index 8e91fe6eb123a..a094c0bcf53a8 100644 --- a/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_stl_vector.py +++ b/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_stl_vector.py @@ -8,8 +8,10 @@ # For the list of contributors see $ROOTSYS/README/CREDITS. # ################################################################################ +import sys + from . import pythonization -from ._rvec import add_array_interface_property +from ._rvec import _array_interface_dtype_map def _data_vec_char(self): @@ -24,6 +26,35 @@ def _data_vec_char(self): self.pop_back() return d + +def _get_array_interface(self): + import ROOT + + value_type = getattr(type(self), "value_type", None) + dtype_numpy = _array_interface_dtype_map.get(value_type) + if dtype_numpy is not None: + dtype_size = ROOT._cppyy.sizeof(value_type) + endianness = "<" if sys.byteorder == "little" else ">" + size = self.size() + # Numpy breaks for data pointer of 0 even though the array is empty. + # We set the pointer to 1 but the value itself is arbitrary and never accessed. + if self.empty(): + pointer = 1 + else: + pointer = ROOT._cppyy.ll.addressof(self.data()) + return { + "shape": (size,), + "typestr": "{}{}{}".format(endianness, dtype_numpy, dtype_size), + "version": 3, + "data": (pointer, False), + } + + +def _add_array_interface_property(klass): + value_type = getattr(klass, "value_type", None) + if value_type in _array_interface_dtype_map: + klass.__array_interface__ = property(_get_array_interface) + @pythonization("vector<", ns="std", is_prefix=True) def pythonize_stl_vector(klass, name): # Parameters: @@ -31,8 +62,7 @@ def pythonize_stl_vector(klass, name): # name: string containing the name of the class # Add numpy array interface - # NOTE: The pythonization is reused from ROOT::VecOps::RVec - add_array_interface_property(klass, name) + _add_array_interface_property(klass) # Inject custom vector::data() value_type = getattr(klass, 'value_type', None)