Skip to content

Commit ae792d5

Browse files
ericmehlericmehl
committed
VectorTypedData : Multidimensional buffer protocol
1 parent 80c3559 commit ae792d5

File tree

3 files changed

+215
-21
lines changed

3 files changed

+215
-21
lines changed

‎include/IECorePython/GeometricTypedDataBinding.inl‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -262,6 +262,7 @@ class GeometricVectorTypedDataFunctions : ThisBinder
262262
"\nor any other python built-in type that is convertible to it. Alternatively accepts the size of the new vector.") \
263263
.def("getInterpretation", &ThisClass::getInterpretation, "Returns the geometric interpretation of this data.") \
264264
.def("setInterpretation", &ThisClass::setInterpretation, "Sets the geometric interpretation of this data.") \
265+
BIND_BUFFER_PROTOCOL_METHODS \
265266
; \
266267
} \
267268

‎src/IECorePython/VectorTypedDataBinding.cpp‎

Lines changed: 109 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -90,6 +90,24 @@ template<> struct PythonType<unsigned short>{ static const char *value(){ return
9090
template<> struct PythonType<int64_t>{ static const char *value(){ return "q"; } };
9191
template<> struct PythonType<uint64_t>{ static const char *value(){ return "Q"; } };
9292

93+
template<typename T>
94+
constexpr std::pair<const char *, Py_ssize_t> typeInfo()
95+
{
96+
if constexpr(
97+
IECore::TypeTraits::IsMatrix<T>::value ||
98+
IECore::TypeTraits::IsColor<T>::value ||
99+
IECore::TypeTraits::IsQuat<T>::value ||
100+
IECore::TypeTraits::IsVec<T>::value
101+
)
102+
{
103+
return { PythonType<typename T::BaseType>::value(), sizeof( T::BaseType ) };
104+
}
105+
else
106+
{
107+
return { PythonType<T>::value(), sizeof( T ) };
108+
}
109+
}
110+
93111
} // namespace
94112

95113
namespace IECorePython
@@ -209,20 +227,95 @@ int Buffer::getBuffer( PyObject *object, Py_buffer *view, int flags )
209227
if constexpr( TypeTraits::HasBaseType<DataType>::value && TypeTraits::IsNumericBasedVectorTypedData<DataType>::value )
210228
{
211229
using ElementType = typename DataType::ValueType::value_type;
212-
if constexpr( std::is_arithmetic_v<ElementType> )
230+
const auto [format, itemSize] = typeInfo<ElementType>();
231+
232+
int ndim = 1;
233+
Py_ssize_t *shape = NULL;
234+
Py_ssize_t *strides = NULL;
235+
if constexpr( IECore::TypeTraits::IsMatrix44<ElementType>::value )
236+
{
237+
ndim = 3;
238+
if( ( flags & PyBUF_ND ) == PyBUF_ND )
239+
{
240+
shape = new Py_ssize_t[3]{ (Py_ssize_t)bufferData->readable().size(), 4, 4 };
241+
}
242+
if( ( flags & PyBUF_STRIDES ) == PyBUF_STRIDES )
243+
{
244+
strides = new Py_ssize_t[3]{ 16 * itemSize, 4 * itemSize, itemSize };
245+
}
246+
}
247+
else if constexpr( IECore::TypeTraits::IsMatrix33<ElementType>::value )
248+
{
249+
ndim = 3;
250+
if( ( flags & PyBUF_ND ) == PyBUF_ND )
251+
{
252+
shape = new Py_ssize_t[3]{ (Py_ssize_t)bufferData->readable().size(), 3, 3 };
253+
}
254+
if( ( flags & PyBUF_STRIDES ) == PyBUF_STRIDES )
255+
{
256+
strides = new Py_ssize_t[3]{ 9 * itemSize, 3 * itemSize, itemSize };
257+
}
258+
}
259+
else if constexpr( IECore::TypeTraits::IsQuat<ElementType>::value )
213260
{
214-
view->obj = Py_XNewRef( object );
215-
view->readonly = !self->isWritable();
216-
view->buf = view->readonly ? (void*)bufferData->baseReadable() : (void*)bufferData->baseWritable();
217-
view->len = bufferData->readable().size() * sizeof( ElementType );
218-
view->itemsize = sizeof( ElementType );
219-
view->format = ( ( flags & PyBUF_FORMAT ) == PyBUF_FORMAT ) ? const_cast<char *>( PythonType<ElementType>::value() ) : NULL;
220-
view->ndim = 1;
221-
view->internal = NULL;
222-
view->shape = ( ( flags & PyBUF_ND ) == PyBUF_ND ) ? (Py_ssize_t*)( new Py_ssize_t( bufferData->readable().size() ) ) : NULL;
223-
view->strides = ( ( flags & PyBUF_STRIDES ) == PyBUF_STRIDES ) ? &( view->itemsize ) : NULL;
224-
view->suboffsets = NULL;
261+
ndim = 2;
262+
if( ( flags & PyBUF_ND ) == PyBUF_ND )
263+
{
264+
shape = new Py_ssize_t[2]{ (Py_ssize_t)bufferData->readable().size(), 4 };
265+
}
266+
if( ( flags & PyBUF_STRIDES ) == PyBUF_STRIDES )
267+
{
268+
strides = new Py_ssize_t[2]{ 4 * itemSize, itemSize };
269+
}
225270
}
271+
else if constexpr( IECore::TypeTraits::IsColor3<ElementType>::value )
272+
{
273+
// `Color3` doesn't have `dimensions()`
274+
ndim = 2;
275+
if( ( flags & PyBUF_ND ) == PyBUF_ND )
276+
{
277+
shape = new Py_ssize_t[2]{ (Py_ssize_t)bufferData->readable().size(), 3 };
278+
}
279+
if( ( flags & PyBUF_STRIDES ) == PyBUF_STRIDES )
280+
{
281+
strides = new Py_ssize_t[2]{ 3 * itemSize, itemSize };
282+
}
283+
}
284+
else if constexpr( IECore::TypeTraits::IsColor4<ElementType>::value || IECore::TypeTraits::IsVec<ElementType>::value )
285+
{
286+
ndim = 2;
287+
if( ( flags & PyBUF_ND ) == PyBUF_ND )
288+
{
289+
shape = new Py_ssize_t[2]{ (Py_ssize_t)bufferData->readable().size(), ElementType::dimensions() };
290+
}
291+
if( ( flags & PyBUF_STRIDES ) == PyBUF_STRIDES )
292+
{
293+
strides = new Py_ssize_t[2]{ ElementType::dimensions() * itemSize, itemSize };
294+
}
295+
}
296+
else
297+
{
298+
shape = new Py_ssize_t{ (Py_ssize_t)bufferData->readable().size() };
299+
}
300+
301+
Py_ssize_t shapeProduct = 1;
302+
Py_ssize_t *idx = shape;
303+
for( int i = 0; i < ndim; ++i, ++idx )
304+
{
305+
shapeProduct *= *idx;
306+
}
307+
308+
view->obj = Py_XNewRef( object );
309+
view->readonly = !self->isWritable();
310+
view->buf = view->readonly ? (void*)bufferData->baseReadable() : (void*)bufferData->baseWritable();
311+
view->len = view->itemsize * shapeProduct;
312+
view->itemsize = itemSize;
313+
view->format = ( ( flags & PyBUF_FORMAT ) == PyBUF_FORMAT ) ? const_cast<char *>( format ) : NULL;
314+
view->ndim = ndim;
315+
view->internal = NULL;
316+
view->shape = shape;
317+
view->strides = strides;
318+
view->suboffsets = NULL;
226319
}
227320
}
228321
);
@@ -243,6 +336,10 @@ void Buffer::releaseBuffer( PyObject *object, Py_buffer *view )
243336
{
244337
delete view->shape;
245338
}
339+
if( view->strides != NULL )
340+
{
341+
delete view->strides;
342+
}
246343
// Python takes care of decrementing `object`
247344
}
248345

‎test/IECore/VectorData.py‎

Lines changed: 105 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1369,6 +1369,18 @@ def test( self ) :
13691369

13701370
class TestBufferProtocol( unittest.TestCase ) :
13711371

1372+
def __assertMemoryBufferProperties( self, pythonBuffer, cortexBuffer, elementFormat, ndim, shape, strides ) :
1373+
1374+
self.assertIs( pythonBuffer.obj, cortexBuffer )
1375+
self.assertTrue( pythonBuffer.readonly )
1376+
self.assertEqual( pythonBuffer.format, elementFormat )
1377+
self.assertEqual( pythonBuffer.ndim, ndim )
1378+
self.assertEqual( pythonBuffer.shape, shape )
1379+
self.assertEqual( pythonBuffer.itemsize, struct.calcsize( elementFormat ) )
1380+
self.assertEqual( pythonBuffer.strides, strides )
1381+
self.assertTrue( pythonBuffer.c_contiguous )
1382+
self.assertTrue( pythonBuffer.contiguous )
1383+
13721384
def testSimpleTypes( self ) :
13731385

13741386
for elementType, vectorType, elementFormat in [
@@ -1393,15 +1405,99 @@ def testSimpleTypes( self ) :
13931405

13941406
# `memoryview` returns `int` for C `char` / Python `chr` types. Cast to `chr`
13951407
self.assertEqual( list( m ) if elementType != chr else [ chr(i) for i in m ], list( v ) )
1396-
self.assertIs( m.obj, b )
1397-
self.assertTrue( m.readonly )
1398-
self.assertEqual( m.format, elementFormat )
1399-
self.assertEqual( m.ndim, 1 )
1400-
self.assertEqual( m.shape, ( len( v ), ) )
1401-
self.assertEqual( m.itemsize, struct.calcsize( elementFormat ) )
1402-
self.assertEqual( m.strides, ( m.itemsize, ) )
1403-
self.assertTrue( m.c_contiguous )
1404-
self.assertTrue( m.contiguous )
1408+
self.__assertMemoryBufferProperties( m, b, elementFormat, 1, ( len( v ), ), ( m.itemsize, ) )
1409+
1410+
def testMatrixTypes( self ) :
1411+
1412+
for elementType, vectorType, matrixSize, elementFormat in [
1413+
( imath.M44f, IECore.M44fVectorData, 4, "f" ),
1414+
( imath.M44d, IECore.M44dVectorData, 4, "d" ),
1415+
( imath.M33f, IECore.M33fVectorData, 3, "f" ),
1416+
( imath.M33d, IECore.M33dVectorData, 3, "d" ),
1417+
] :
1418+
with self.subTest( elementType = elementType, vectorType = vectorType ) :
1419+
v = vectorType(
1420+
[
1421+
elementType( *list( range( 0, matrixSize * matrixSize ) ) ),
1422+
elementType( *list( range( 1, 1 + matrixSize * matrixSize ) ) ),
1423+
elementType( *list( range( 2, 2 + matrixSize * matrixSize ) ) ),
1424+
]
1425+
)
1426+
1427+
b = v.asReadOnlyBuffer()
1428+
m = memoryview( b )
1429+
1430+
self.assertIsNotNone( m )
1431+
for i in range( 0, len( v ) ) :
1432+
for j in range( 0, matrixSize ) :
1433+
for k in range( 0, matrixSize ) :
1434+
self.assertEqual( m.tolist()[i][j][k], v[i][j][k] )
1435+
1436+
self.__assertMemoryBufferProperties(
1437+
m,
1438+
b,
1439+
elementFormat,
1440+
3,
1441+
( len( v ), matrixSize, matrixSize ),
1442+
( matrixSize * matrixSize * m.itemsize, matrixSize * m.itemsize, m.itemsize )
1443+
)
1444+
1445+
def testQuat( self ) :
1446+
1447+
for elementType, vectorType, dimensions, elementFormat in [
1448+
( imath.Quatf, IECore.QuatfVectorData, 4, "f" ),
1449+
( imath.Quatd, IECore.QuatdVectorData, 4, "d" ),
1450+
] :
1451+
with self.subTest( elementType = elementType, vectorType = vectorType ) :
1452+
v = vectorType(
1453+
[
1454+
elementType( *list( range( 0, dimensions ) ) ),
1455+
elementType( *list( range( 1, 1 + dimensions ) ) ),
1456+
elementType( *list( range( 2, 2 + dimensions ) ) ),
1457+
]
1458+
)
1459+
1460+
b = v.asReadOnlyBuffer()
1461+
m = memoryview( b )
1462+
1463+
self.assertIsNotNone( m )
1464+
for i in range( 0, len( v ) ) :
1465+
mList = m.tolist()
1466+
self.assertEqual( elementType( mList[i][0], mList[i][1], mList[i][2], mList[i][3] ), v[i] )
1467+
1468+
self.__assertMemoryBufferProperties( m, b, elementFormat, 2, ( len( v ), dimensions ), ( dimensions * m.itemsize, m.itemsize ) )
1469+
1470+
def testTwoDimensionalBuffers( self ) :
1471+
1472+
for elementType, vectorType, dimensions, elementFormat in [
1473+
( imath.Color3f, IECore.Color3fVectorData, 3, "f" ),
1474+
( imath.Color4f, IECore.Color4fVectorData, 4, "f" ),
1475+
( imath.V2f, IECore.V2fVectorData, 2, "f" ),
1476+
( imath.V2d, IECore.V2dVectorData, 2, "d" ),
1477+
( imath.V2i, IECore.V2iVectorData, 2, "i" ),
1478+
( imath.V3f, IECore.V3fVectorData, 3, "f" ),
1479+
( imath.V3d, IECore.V3dVectorData, 3, "d" ),
1480+
( imath.V3i, IECore.V3iVectorData, 3, "i" ),
1481+
] :
1482+
with self.subTest( elementType = elementType, vectorType = vectorType ) :
1483+
v = vectorType(
1484+
[
1485+
elementType( *list( range( 0, dimensions ) ) ),
1486+
elementType( *list( range( 1, 1 + dimensions ) ) ),
1487+
elementType( *list( range( 2, 2 + dimensions ) ) ),
1488+
]
1489+
)
1490+
1491+
b = v.asReadOnlyBuffer()
1492+
m = memoryview( b )
1493+
1494+
self.assertIsNotNone( m )
1495+
for i in range( 0, len( v ) ) :
1496+
for j in range( 0, dimensions ) :
1497+
self.assertEqual( m.tolist()[i][j], v[i][j] )
1498+
1499+
self.__assertMemoryBufferProperties( m, b, elementFormat, 2, ( len( v ), dimensions ), ( dimensions * m.itemsize, m.itemsize ) )
1500+
14051501

14061502
def testReadOnlyBuffer( self ) :
14071503

0 commit comments

Comments
 (0)