cimport libav as lib
from libc.string cimport memcpy
from av.error cimport stash_exception
ctypedef int64_t (*seek_func_t)(void *opaque, int64_t offset, int whence) noexcept nogil
cdef class PyIOFile:
def __cinit__(self, file, buffer_size, writeable=None):
self.file = file
cdef seek_func_t seek_func = NULL
readable = getattr(self.file, "readable", None)
writable = getattr(self.file, "writable", None)
seekable = getattr(self.file, "seekable", None)
self.fread = getattr(self.file, "read", None)
self.fwrite = getattr(self.file, "write", None)
self.fseek = getattr(self.file, "seek", None)
self.ftell = getattr(self.file, "tell", None)
self.fclose = getattr(self.file, "close", None)
# To be seekable the file object must have `seek` and `tell` methods.
# If it also has a `seekable` method, it must return True.
if (
self.fseek is not None
and self.ftell is not None
and (seekable is None or seekable())
):
seek_func = pyio_seek
if writeable is None:
writeable = self.fwrite is not None
if writeable:
if self.fwrite is None or (writable is not None and not writable()):
raise ValueError("File object has no write() method, or writable() returned False.")
else:
if self.fread is None or (readable is not None and not readable()):
raise ValueError("File object has no read() method, or readable() returned False.")
self.pos = 0
self.pos_is_valid = True
# This is effectively the maximum size of reads.
self.buffer = <unsigned char*>lib.av_malloc(buffer_size)
self.iocontext = lib.avio_alloc_context(
self.buffer,
buffer_size,
writeable,
<void*>self, # User data.
pyio_read,
pyio_write,
seek_func
)
if seek_func:
self.iocontext.seekable = lib.AVIO_SEEKABLE_NORMAL
self.iocontext.max_packet_size = buffer_size
def __dealloc__(self):
with nogil:
# FFmpeg will not release custom input, so it's up to us to free it.
# Do not touch our original buffer as it may have been freed and replaced.
if self.iocontext:
lib.av_freep(&self.iocontext.buffer)
lib.av_freep(&self.iocontext)
# We likely errored badly if we got here, and so are still
# responsible for our buffer.
else:
lib.av_freep(&self.buffer)
cdef int pyio_read(void *opaque, uint8_t *buf, int buf_size) noexcept nogil:
with gil:
return pyio_read_gil(opaque, buf, buf_size)
cdef int pyio_read_gil(void *opaque, uint8_t *buf, int buf_size) noexcept:
cdef PyIOFile self
cdef bytes res
try:
self = <PyIOFile>opaque
res = self.fread(buf_size)
memcpy(buf, <void*><char*>res, len(res))
self.pos += len(res)
if not res:
return lib.AVERROR_EOF
return len(res)
except Exception:
return stash_exception()
cdef int pyio_write(void *opaque, const uint8_t *buf, int buf_size) noexcept nogil:
with gil:
return pyio_write_gil(opaque, buf, buf_size)
cdef int pyio_write_gil(void *opaque, const uint8_t *buf, int buf_size) noexcept:
cdef PyIOFile self
cdef bytes bytes_to_write
cdef int bytes_written
try:
self = <PyIOFile>opaque
bytes_to_write = buf[:buf_size]
ret_value = self.fwrite(bytes_to_write)
bytes_written = ret_value if isinstance(ret_value, int) else buf_size
self.pos += bytes_written
return bytes_written
except Exception:
return stash_exception()
cdef int64_t pyio_seek(void *opaque, int64_t offset, int whence) noexcept nogil:
# Seek takes the standard flags, but also a ad-hoc one which means that
# the library wants to know how large the file is. We are generally
# allowed to ignore this.
if whence == lib.AVSEEK_SIZE:
return -1
with gil:
return pyio_seek_gil(opaque, offset, whence)
cdef int64_t pyio_seek_gil(void *opaque, int64_t offset, int whence):
cdef PyIOFile self
try:
self = <PyIOFile>opaque
res = self.fseek(offset, whence)
# Track the position for the user.
if whence == 0:
self.pos = offset
elif whence == 1:
self.pos += offset
else:
self.pos_is_valid = False
if res is None:
if self.pos_is_valid:
res = self.pos
else:
res = self.ftell()
return res
except Exception:
return stash_exception()
cdef int pyio_close_gil(lib.AVIOContext *pb):
try:
return lib.avio_close(pb)
except Exception:
stash_exception()
cdef int pyio_close_custom_gil(lib.AVIOContext *pb):
cdef PyIOFile self
try:
self = <PyIOFile>pb.opaque
# Flush bytes in the AVIOContext buffers to the custom I/O
lib.avio_flush(pb)
if self.fclose is not None:
self.fclose()
return 0
except Exception:
stash_exception()