101 lines
2.8 KiB
Python
101 lines
2.8 KiB
Python
|
import cython
|
||
|
|
from cython.cimports.cpython import PyBUF_WRITABLE, PyBuffer_FillInfo
|
||
|
|
from cython.cimports.cpython.buffer import (
|
||
|
|
PyBUF_SIMPLE,
|
||
|
|
PyBuffer_Release,
|
||
|
|
PyObject_CheckBuffer,
|
||
|
|
PyObject_GetBuffer,
|
||
|
|
)
|
||
|
|
from cython.cimports.libc.string import memcpy
|
||
|
|
|
||
|
|
|
||
|
|
@cython.final
|
||
|
|
@cython.cclass
|
||
|
|
class ByteSource:
|
||
|
|
def __cinit__(self, owner):
|
||
|
|
self.owner = owner
|
||
|
|
|
||
|
|
try:
|
||
|
|
self.ptr = owner
|
||
|
|
except TypeError:
|
||
|
|
pass
|
||
|
|
else:
|
||
|
|
self.length = len(owner)
|
||
|
|
return
|
||
|
|
|
||
|
|
if PyObject_CheckBuffer(owner):
|
||
|
|
# Can very likely use PyBUF_ND instead of PyBUF_SIMPLE
|
||
|
|
res = PyObject_GetBuffer(owner, cython.address(self.view), PyBUF_SIMPLE)
|
||
|
|
if not res:
|
||
|
|
self.has_view = True
|
||
|
|
self.ptr = cython.cast(cython.p_uchar, self.view.buf)
|
||
|
|
self.length = self.view.len
|
||
|
|
return
|
||
|
|
|
||
|
|
raise TypeError("expected bytes, bytearray or memoryview")
|
||
|
|
|
||
|
|
def __dealloc__(self):
|
||
|
|
if self.has_view:
|
||
|
|
PyBuffer_Release(cython.address(self.view))
|
||
|
|
|
||
|
|
|
||
|
|
@cython.cfunc
|
||
|
|
def bytesource(obj, allow_none: cython.bint = False) -> ByteSource | None:
|
||
|
|
if allow_none and obj is None:
|
||
|
|
return None
|
||
|
|
elif isinstance(obj, ByteSource):
|
||
|
|
return obj
|
||
|
|
else:
|
||
|
|
return ByteSource(obj)
|
||
|
|
|
||
|
|
|
||
|
|
@cython.cclass
|
||
|
|
class Buffer:
|
||
|
|
"""A base class for PyAV objects which support the buffer protocol, such
|
||
|
|
as :class:`.Packet` and :class:`.Plane`.
|
||
|
|
|
||
|
|
"""
|
||
|
|
|
||
|
|
@cython.cfunc
|
||
|
|
def _buffer_size(self) -> cython.size_t:
|
||
|
|
return 0
|
||
|
|
|
||
|
|
def _buffer_ptr(self) -> cython.p_void:
|
||
|
|
return cython.NULL
|
||
|
|
|
||
|
|
def _buffer_writable(self) -> cython.bint:
|
||
|
|
return True
|
||
|
|
|
||
|
|
def __getbuffer__(self, view: cython.pointer[Py_buffer], flags: cython.int):
|
||
|
|
if flags & PyBUF_WRITABLE and not self._buffer_writable():
|
||
|
|
raise ValueError("buffer is not writable")
|
||
|
|
|
||
|
|
PyBuffer_FillInfo(view, self, self._buffer_ptr(), self._buffer_size(), 0, flags)
|
||
|
|
|
||
|
|
@property
|
||
|
|
def buffer_size(self):
|
||
|
|
return self._buffer_size()
|
||
|
|
|
||
|
|
@property
|
||
|
|
def buffer_ptr(self):
|
||
|
|
"""The memory address of the buffer."""
|
||
|
|
return cython.cast(cython.size_t, self._buffer_ptr())
|
||
|
|
|
||
|
|
def update(self, input):
|
||
|
|
"""Replace the data in this object with the given buffer.
|
||
|
|
|
||
|
|
Accepts anything that supports the `buffer protocol <https://docs.python.org/3/c-api/buffer.html>`_,
|
||
|
|
e.g. bytes, NumPy arrays, other :class:`Buffer` objects, etc..
|
||
|
|
|
||
|
|
"""
|
||
|
|
if not self._buffer_writable():
|
||
|
|
raise ValueError("buffer is not writable")
|
||
|
|
|
||
|
|
source: ByteSource = bytesource(input)
|
||
|
|
size: cython.size_t = self._buffer_size()
|
||
|
|
|
||
|
|
if source.length != size:
|
||
|
|
raise ValueError(f"got {source.length} bytes; need {size} bytes")
|
||
|
|
|
||
|
|
memcpy(self._buffer_ptr(), source.ptr, size)
|