rk_pve/mp4/decoder.py

# RozK

from . import libav
from .packet import Packet
from .frame import Frame

class Decoder:
    __slots__ = '_context'

    def __init__(self, stream):
        self._context = libav.codec_alloc_context(stream.codec)
        if not self._context:
            raise MemoryError
        errcode = libav.codec_parameters_to_context(self._context, stream.parameters)
        if errcode < 0:
            libav.codec_free_context(self._context)
            raise Exception("Failed to set context parameters")
        errcode = libav.codec_open(self._context, stream.codec)
        if errcode < 0:
            libav.codec_free_context(self._context)
            raise Exception("Failed to open codec context")

    def __del__(self):
        if self._context:
            libav.codec_free_context(self._context)

    def _receive(self):
        frames = []
        while True:
            frame = Frame()
            errcode = libav.codec_receive_frame(self._context, frame)
            if errcode in (libav.AVERROR_EOF, libav.AVERROR_EAGAIN):
                break
            elif errcode < 0:
                errstring = libav.strerror(errcode)
                raise Exception(f"Failed to receive frame: {errstring}")
            frames.append(frame)
        return frames

    def decode(self, packet):
        if not self._context:
            return None
        errcode = libav.codec_send_packet(self._context, packet)
        if errcode < 0:
            errstring = libav.strerror(errcode)
            raise Exception(f"Failed to send packet: {errstring}")
        return self._receive()

    def flush(self):
        return self.decode(None)