VectorTypedData : Support Python buffer protocol

Author: ericmehl

SConstruct

@@ -93,10 +93,13 @@ o.Add(
 # https://developercommunity.visualstudio.com/content/problem/756694/including-windowsh-and-boostinterprocess-headers-l.html
 # /DBOOST_ALL_NO_LIB is needed to find Boost when it is built without
 # verbose system information added to file and directory names.
+# /DZc:strictStrings- fixes a compilation error when setting the `Py_buffer.format` member in `VectorTypedDataBinding.inl`
+# `getBuffer()` method. Python declares that member as `char *` but MSVC requires `const char *` for string literals.
+# Disabling strict strings relaxes that requirement.
 o.Add(
 	"CXXFLAGS",
 	"The extra flags to pass to the C++ compiler during compilation.",
-	[ "-pipe", "-Wall", "-Wextra", "-Wsuggest-override" ] if Environment()["PLATFORM"] != "win32" else [ "/permissive-", "/D_USE_MATH_DEFINES", "/Zc:externC-", "/DBOOST_ALL_NO_LIB" ],
+	[ "-pipe", "-Wall", "-Wextra", "-Wsuggest-override" ] if Environment()["PLATFORM"] != "win32" else [ "/permissive-", "/D_USE_MATH_DEFINES", "/Zc:externC-", "/DBOOST_ALL_NO_LIB", "/Zc:strictStrings-" ],
 )
 
 o.Add(

include/IECorePython/VectorTypedDataBinding.h

@@ -37,8 +37,35 @@
 
 #include "IECorePython/Export.h"
 
+#include "IECore/Data.h"
+#include "IECore/RefCounted.h"
+
+
 namespace IECorePython
 {
+
+class Buffer : public IECore::RefCounted
+{
+	public :
+		IE_CORE_DECLAREMEMBERPTR( Buffer );
+
+		Buffer( IECore::Data *data, const bool writable );
+		~Buffer() override;
+
+		IECore::DataPtr asData() const;
+
+		bool isWritable() const;
+
+		static int getBuffer( PyObject *object, Py_buffer *view, int flags );
+		static void releaseBuffer( PyObject *object, Py_buffer *view );
+
+	private :
+		IECore::DataPtr m_data;
+		const bool m_writable;
+};
+
+IE_CORE_DECLAREPTR( Buffer )
+
 extern IECOREPYTHON_API void bindAllVectorTypedData();
 }
 

include/IECorePython/VectorTypedDataBinding.inl

@@ -36,14 +36,21 @@
 #define IECOREPYTHON_VECTORTYPEDDATABINDING_INL
 
 #include "boost/python.hpp"
+#include "Python.h"
 
 #include "IECorePython/IECoreBinding.h"
 #include "IECorePython/RunTimeTypedBinding.h"
+#include "IECorePython/VectorTypedDataBinding.h"
+
+#include "IECore/MessageHandler.h"
+#include "IECore/TypeTraits.h"
 
 #include "boost/numeric/conversion/cast.hpp"
+
 #include "boost/python/suite/indexing/container_utils.hpp"
 
 #include <sstream>
+#include <typeindex>
 
 namespace IECorePython
 {
@@ -546,6 +553,16 @@ class VectorTypedDataFunctions
 			return &x.readable() == &y.readable();
 		}
 
+		static IECorePython::BufferPtr asReadOnlyBuffer( ThisClass &x )
+		{
+			return new Buffer( &x, /* writable = */ false );
+		}
+
+		static IECorePython::BufferPtr asReadWriteBuffer( ThisClass &x )
+		{
+			return new Buffer( &x, /* writable = */ true );
+		}
+
 
 	protected:
 		/*
@@ -700,7 +717,7 @@ std::string str<IECore::TypedData<std::vector<TYPE> > >( IECore::TypedData<std::
 #define BASIC_VECTOR_BINDING(ThisClass, Tname)																	\
 		typedef VectorTypedDataFunctions< ThisClass > ThisBinder;												\
 																													\
-		RunTimeTypedClass<ThisClass>(																							\
+		RunTimeTypedClass<ThisClass>(																				\
 			Tname "-type vector class derived from Data class.\n"													\
 			"This class behaves like the native python lists, except that it only accepts " Tname " values.\n"		\
 			"The copy constructor accepts another instance of this class or a python list containing " Tname		\
@@ -742,6 +759,10 @@ std::string str<IECore::TypedData<std::vector<TYPE> > >( IECore::TypedData<std::
 			;																						\
 		}
 
+#define BIND_BUFFER_PROTOCOL_METHODS	\
+				.def("asReadOnlyBuffer", &ThisBinder::asReadOnlyBuffer, "Returns a read-only Buffer object that can be passed to Python objects that support the buffer protocol." )	\
+				.def("asReadWriteBuffer", &ThisBinder::asReadWriteBuffer, "Returns a writable Buffer object that can be passed to Python objects that support the buffer protocol." )	\
+
 // bind a VectorTypedData class that supports simple Math operators (+=, -= and *=)
 #define BIND_SIMPLE_OPERATED_VECTOR_TYPEDDATA(T, Tname)									\
 		{																							\
@@ -755,6 +776,7 @@ std::string str<IECore::TypedData<std::vector<TYPE> > >( IECore::TypedData<std::
 				.def("__imul__", &ThisBinder::imul, "inplace multiplication (s *= v) : accepts another vector of the same type or a single " Tname)		\
 				.def("__cmp__", &ThisBinder::invalidOperator, "Raises an exception. This vector type does not support comparison operators.")		\
 				.def("toString", &ThisBinder::toString, "Returns a string with a copy of the bytes in the vector.")\
+				BIND_BUFFER_PROTOCOL_METHODS														\
 			;																						\
 		}
 
@@ -775,6 +797,7 @@ std::string str<IECore::TypedData<std::vector<TYPE> > >( IECore::TypedData<std::
 				.def("__itruediv__", &ThisBinder::idiv, "inplace division (s /= v) : accepts another vector of the same type or a single " Tname)			\
 				.def("__cmp__", &ThisBinder::invalidOperator, "Raises an exception. This vector type does not support comparison operators.")		\
 				.def("toString", &ThisBinder::toString, "Returns a string with a copy of the bytes in the vector.")\
+				BIND_BUFFER_PROTOCOL_METHODS																											\
 			;																						\
 		}
 
@@ -799,6 +822,7 @@ std::string str<IECore::TypedData<std::vector<TYPE> > >( IECore::TypedData<std::
 				.def("__gt__", &ThisBinder::gt, "The comparison is element-wise, like a string comparison. \n")	\
 				.def("__ge__", &ThisBinder::ge, "The comparison is element-wise, like a string comparison. \n")	\
 				.def("toString", &ThisBinder::toString, "Returns a string with a copy of the bytes in the vector.")\
+				BIND_BUFFER_PROTOCOL_METHODS																		\
 			; \
 		}
 

src/IECorePython/VectorTypedDataBinding.cpp

@@ -47,6 +47,7 @@
 #include "IECorePython/RunTimeTypedBinding.h"
 #include "IECorePython/VectorTypedDataBinding.inl"
 
+#include "IECore/DataAlgo.h"
 #include "IECore/Export.h"
 #include "IECore/VectorTypedData.h"
 
@@ -68,6 +69,29 @@ using namespace boost::python;
 using namespace Imath;
 using namespace IECore;
 
+namespace
+{
+
+template<typename T>
+struct PythonType
+{
+	static const char *value();
+};
+
+template<> struct PythonType<half>{ static const char *value(){ return "e"; } };
+template<> struct PythonType<float>{ static const char *value(){ return "f"; } };
+template<> struct PythonType<double>{ static const char *value(){ return "d"; } };
+template<> struct PythonType<int>{ static const char *value(){ return "i"; } };
+template<> struct PythonType<unsigned int>{ static const char *value(){ return "I"; } };
+template<> struct PythonType<char>{ static const char *value(){ return "b"; } };
+template<> struct PythonType<unsigned char>{ static const char *value(){ return "B"; } };
+template<> struct PythonType<short>{ static const char *value(){ return "h"; } };
+template<> struct PythonType<unsigned short>{ static const char *value(){ return "H"; } };
+template<> struct PythonType<int64_t>{ static const char *value(){ return "q"; } };
+template<> struct PythonType<uint64_t>{ static const char *value(){ return "Q"; } };
+
+}  // namespace
+
 namespace IECorePython
 {
 
@@ -124,6 +148,109 @@ std::string str<BoolVectorData>( BoolVectorData &x )
 	return s.str();
 }
 
+Buffer::Buffer( Data *data, const bool writable ) : m_data( data->copy() ), m_writable( writable )
+{
+
+}
+
+Buffer::~Buffer()
+{
+
+}
+
+DataPtr Buffer::asData() const
+{
+	return m_data->copy();
+}
+
+bool Buffer::isWritable() const
+{
+	return m_writable;
+}
+
+int Buffer::getBuffer( PyObject *object, Py_buffer *view, int flags )
+{
+	// This method is a customized variation on Python's `PyBuffer_FillInfo` to suit our needs.
+	if( view == NULL ) {
+		PyErr_SetString( PyExc_ValueError, "getBuffer(): view==NULL argument is obsolete" );
+		return -1;
+	}
+
+	if( flags != PyBUF_SIMPLE )
+	{
+		if( flags == PyBUF_READ || flags == PyBUF_WRITE )
+		{
+			PyErr_BadInternalCall();
+			return -1;
+		}
+	}
+
+	IECorePython::BufferPtr self = boost::python::extract<IECorePython::BufferPtr>( object );
+	if( !self )
+	{
+		/// \todo reword this
+		PyErr_SetString( PyExc_ValueError, "getBuffer(): Buffer type does not match expected type." );
+		return -1;
+	}
+
+	if( ( flags | PyBUF_WRITABLE ) == PyBUF_WRITABLE && !self->isWritable() )
+	{
+		return -1;
+	}
+
+	try
+	{
+		dispatch(
+			self->m_data.get(),
+			[&flags, &view, &object, &self]( auto *bufferData ) -> void
+			{
+				using DataType = typename std::remove_const_t< std::remove_pointer_t< decltype( bufferData ) > >;
+
+				if constexpr( TypeTraits::HasBaseType<DataType>::value && TypeTraits::IsNumericBasedVectorTypedData<DataType>::value )
+				{
+					using ElementType = typename DataType::ValueType::value_type;
+					if constexpr( std::is_arithmetic_v<ElementType> )
+					{
+						view->obj = Py_XNewRef( object );
+						view->readonly = !self->isWritable();
+						view->buf = view->readonly ? (void*)bufferData->baseReadable() : (void*)bufferData->baseWritable();
+						view->len = bufferData->readable().size() * sizeof( ElementType );
+						view->itemsize = sizeof( ElementType );
+						view->format = ( ( flags & PyBUF_FORMAT ) == PyBUF_FORMAT ) ? const_cast<char *>( PythonType<ElementType>::value() ) : NULL;
+						view->ndim = 1;
+						view->internal = NULL;
+						view->shape = ( ( flags & PyBUF_ND ) == PyBUF_ND ) ? (Py_ssize_t*)( new Py_ssize_t( bufferData->readable().size() ) ) : NULL;
+						view->strides = ( ( flags & PyBUF_STRIDES ) == PyBUF_STRIDES ) ? &( view->itemsize ) : NULL;
+						view->suboffsets = NULL;
+					}
+				}
+			}
+		);
+	}
+	catch( Exception &e )
+	{
+		view->obj = NULL;
+		PyErr_SetString( PyExc_BufferError, e.what() );
+		return -1;
+	}
+
+	return 0;
+}
+
+void Buffer::releaseBuffer( PyObject *object, Py_buffer *view )
+{
+	if( view->shape != NULL )
+	{
+		delete view->shape;
+	}
+	// Python takes care of decrementing `object`
+}
+
+static PyBufferProcs BufferProtocol = {
+  (getbufferproc)Buffer::getBuffer,
+  (releasebufferproc)Buffer::releaseBuffer,
+};
+
 void bindAllVectorTypedData()
 {
 	// basic types
@@ -189,6 +316,16 @@ void bindAllVectorTypedData()
 	bindImathColorVectorTypedData();
 	bindImathBoxVectorTypedData();
 	bindImathQuatVectorTypedData();
+
+	auto c = RefCountedClass<Buffer, IECore::RefCounted>( "Buffer" )
+		.def( init<Data *, bool>() )
+		.def( "asData", &Buffer::asData )
+		.def( "isWritable", &Buffer::isWritable )
+	;
+	PyTypeObject *o = (PyTypeObject *)c.ptr();
+	o->tp_as_buffer = &BufferProtocol;
+	PyType_Modified( o );
+	
 }
 
 

test/IECore/VectorData.py

@@ -37,6 +37,7 @@
 import math
 import os
 import unittest
+import struct
 import imath
 
 import IECore
@@ -1366,6 +1367,91 @@ def test( self ) :
 		self.assertFalse( a._dataSourceEqual( ca2 ) )
 		self.assertTrue(ca2._dataSourceEqual( ca3 ) )
 
+class TestBufferProtocol( unittest.TestCase ) :
+
+	def testSimpleTypes( self ) :
+
+		for elementType, vectorType, elementFormat in [
+			# It would be nice to test for `IECore.HalfVectorData: "e"`
+			# but that is not supported in our Python.
+			( float, IECore.FloatVectorData, "f" ),
+			( float, IECore.DoubleVectorData, "d" ),
+			( int, IECore.IntVectorData, "i" ),
+			( int, IECore.UIntVectorData, "I" ),
+			( chr, IECore.CharVectorData, "b" ),
+			( int, IECore.UCharVectorData, "B" ),
+			( int, IECore.ShortVectorData, "h" ),
+			( int, IECore.UShortVectorData, "H" ),
+			( int, IECore.Int64VectorData, "q" ),
+			( int, IECore.UInt64VectorData, "Q" ),
+		] :
+			with self.subTest( elementType = elementType, vectorType = vectorType ) :
+				v = vectorType( [ elementType( 1 ), elementType( 2 ), elementType( 3 ) ] )
+
+				b = v.asReadOnlyBuffer()
+				m = memoryview( b )
+
+				# `memoryview` returns `int` for C `char` / Python `chr` types. Cast to `chr`
+				self.assertEqual( list( m ) if elementType != chr else [ chr(i) for i in m ], list( v ) )
+				self.assertIs( m.obj, b )
+				self.assertTrue( m.readonly )
+				self.assertEqual( m.format, elementFormat )
+				self.assertEqual( m.ndim, 1 )
+				self.assertEqual( m.shape, ( len( v ), ) )
+				self.assertEqual( m.itemsize, struct.calcsize( elementFormat ) )
+				self.assertEqual( m.strides, ( m.itemsize, ) )
+				self.assertTrue( m.c_contiguous )
+				self.assertTrue( m.contiguous )
+
+	def testReadOnlyBuffer( self ) :
+
+		v = IECore.FloatVectorData( [ 1, 2, 3 ] )
+		buffer = v.asReadOnlyBuffer()
+		bufferData = buffer.asData()
+		self.assertFalse( buffer.isWritable() )
+		self.assertTrue( v._dataSourceEqual( bufferData ) )
+		self.assertTrue( bufferData._dataSourceEqual( buffer.asData() ) )
+
+		m = memoryview( buffer )
+
+		self.assertTrue( m.readonly )
+		self.assertTrue( v._dataSourceEqual( bufferData ) )
+		self.assertTrue( bufferData._dataSourceEqual( buffer.asData() ) )
+
+		# We can modify `v` without affecting our buffer.
+		v.append( 99 )
+		self.assertEqual( list( m ), [ 1, 2, 3 ] )
+		self.assertFalse( v._dataSourceEqual( bufferData ) )
+		self.assertTrue( bufferData._dataSourceEqual( buffer.asData() ) )
+
+		# We can modify `bufferData`. It's a copy so writing to it results
+		# in a deep copy of the source data.
+		bufferData[0] = 42
+		self.assertEqual( list( bufferData ), [ 42, 2, 3 ] )
+		self.assertFalse( bufferData._dataSourceEqual( buffer.asData() ) )
+		self.assertEqual( list( buffer.asData() ), [ 1, 2, 3 ] )
+
+	def testReadWriteBuffer( self ) :
+
+		v = IECore.FloatVectorData( [ 1, 2, 3 ] )
+		buffer = v.asReadWriteBuffer()
+		bufferData = buffer.asData()
+		self.assertTrue( buffer.isWritable() )
+		self.assertTrue( v._dataSourceEqual( bufferData ) )
+		self.assertTrue( bufferData._dataSourceEqual( buffer.asData() ) )
+
+		# Creating a writable `memoryview` makes `bufferData` unique.
+		m = memoryview( buffer )
+		self.assertFalse( m.readonly )
+		self.assertFalse( v._dataSourceEqual( buffer.asData() ) )
+		self.assertFalse( bufferData._dataSourceEqual( buffer.asData() ) )
+
+		# Modify the memoryview, which is reflected in the buffer.
+		m[0] = 42
+		self.assertEqual( list( m ), [ 42, 2, 3 ] )
+		self.assertEqual( list( buffer.asData() ), list( m ) )
+
+
 if __name__ == "__main__":
     unittest.main()