Combine av/buffer and av/bytesource

Author: WyattBlue

av/buffer.pxd

@@ -1,6 +1,16 @@
+from cpython.buffer cimport Py_buffer
 
-cdef class Buffer:
 
+cdef class ByteSource:
+    cdef object owner
+    cdef bint has_view
+    cdef Py_buffer view
+    cdef unsigned char *ptr
+    cdef size_t length
+
+cdef ByteSource bytesource(object, bint allow_none=*)
+
+cdef class Buffer:
     cdef size_t _buffer_size(self)
     cdef void* _buffer_ptr(self)
     cdef bint _buffer_writable(self)

av/buffer.py

@@ -1,9 +1,53 @@
 import cython
-from cython.cimports.av.bytesource import ByteSource, bytesource
 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.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
@@ -29,7 +73,6 @@ def __getbuffer__(self, view: cython.pointer[Py_buffer], flags: cython.int):
 
     @property
     def buffer_size(self):
-        """The size of the buffer in bytes."""
         return self._buffer_size()
 
     @property

av/bytesource.pxd

@@ -1,7 +1,7 @@
 cimport libav as lib
 from libc.stdint cimport int64_t
 
-from av.bytesource cimport ByteSource
+from av.buffer cimport ByteSource
 from av.codec.codec cimport Codec
 from av.codec.hwaccel cimport HWAccel
 from av.frame cimport Frame

av/bytesource.pyx

@@ -2,7 +2,7 @@
 
 import cython
 from cython.cimports import libav as lib
-from cython.cimports.av.bytesource import ByteSource, bytesource
+from cython.cimports.av.buffer import ByteSource, bytesource
 from cython.cimports.av.codec.codec import Codec, wrap_codec
 from cython.cimports.av.dictionary import _Dictionary
 from cython.cimports.av.error import err_check

av/codec/context.pxd

@@ -2,8 +2,7 @@ from cython.cimports.libc.stdint import uint8_t
 
 cimport libav as lib
 
-from av.buffer cimport Buffer
-from av.bytesource cimport ByteSource
+from av.buffer cimport Buffer, ByteSource
 from av.stream cimport Stream
 
 

av/codec/context.py

@@ -2,8 +2,7 @@
 
 import cython
 from cython.cimports import libav as lib
-from cython.cimports.av.buffer import Buffer
-from cython.cimports.av.bytesource import ByteSource, bytesource
+from cython.cimports.av.buffer import Buffer, ByteSource, bytesource
 from cython.cimports.av.error import err_check
 from cython.cimports.av.opaque import noop_free, opaque_container
 from cython.cimports.av.utils import avrational_to_fraction, to_avrational

av/packet.pxd

@@ -1,6 +1,6 @@
 import cython
 from cython.cimports import libav as lib
-from cython.cimports.av.bytesource import ByteSource, bytesource
+from cython.cimports.av.buffer import ByteSource, bytesource
 from cython.cimports.av.error import err_check
 from cython.cimports.av.packet import Packet
 from cython.cimports.av.subtitles.subtitle import SubtitleProxy, SubtitleSet