Make av/sidedata/encparams pure

av/sidedata/encparams.pyx → av/sidedata/encparams.py

@@ -1,19 +1,20 @@
-cimport libav as lib
-from libc.stdint cimport int32_t, uint8_t
-
 from enum import IntEnum
 
-VideoEncParamsType = IntEnum(
-    "AVVideoEncParamsType",
-    {
-        "NONE":  <int>lib.AV_VIDEO_ENC_PARAMS_NONE,
-        "VP9":   <int>lib.AV_VIDEO_ENC_PARAMS_VP9,
-        "H264":  <int>lib.AV_VIDEO_ENC_PARAMS_H264,
-        "MPEG2": <int>lib.AV_VIDEO_ENC_PARAMS_MPEG2,
-    },
-)
-
-cdef class VideoEncParams(SideData):
+import cython
+from cython.cimports import libav as lib
+from cython.cimports.av.sidedata.sidedata import SideData
+from cython.cimports.libc.stdint import uint8_t
+
+
+class VideoEncParamsType(IntEnum):
+    NONE = lib.AV_VIDEO_ENC_PARAMS_NONE
+    VP9 = lib.AV_VIDEO_ENC_PARAMS_VP9
+    H264 = lib.AV_VIDEO_ENC_PARAMS_H264
+    MPEG2 = lib.AV_VIDEO_ENC_PARAMS_MPEG2
+
+
+@cython.cclass
+class VideoEncParams(SideData):
     def __repr__(self):
         return f"<av.sidedata.VideoEncParams, nb_blocks={self.nb_blocks}, codec_type={self.codec_type}, qp={self.qp}>"
 
@@ -25,29 +26,37 @@ def nb_blocks(self):
         the values of blocks_offset / block_size are unspecified and should not
         be accessed.
         """
-        return (<lib.AVVideoEncParams*> self.ptr.data).nb_blocks
+        return cython.cast(
+            cython.pointer[lib.AVVideoEncParams], self.ptr.data
+        ).nb_blocks
 
     @property
     def blocks_offset(self):
         """
         Offset in bytes from the beginning of this structure at which the array of blocks starts.
         """
-        return (<lib.AVVideoEncParams*> self.ptr.data).blocks_offset
+        return cython.cast(
+            cython.pointer[lib.AVVideoEncParams], self.ptr.data
+        ).blocks_offset
 
     @property
     def block_size(self):
         """
         Size of each block in bytes. May not match sizeof(AVVideoBlockParams).
         """
-        return (<lib.AVVideoEncParams*> self.ptr.data).block_size
+        return cython.cast(
+            cython.pointer[lib.AVVideoEncParams], self.ptr.data
+        ).block_size
 
     @property
     def codec_type(self):
         """
         Type of the parameters (the codec they are used with).
         """
-        cdef lib.AVVideoEncParamsType t = (<lib.AVVideoEncParams*> self.ptr.data).type
-        return VideoEncParamsType(<int>t)
+        t: lib.AVVideoEncParamsType = cython.cast(
+            cython.pointer[lib.AVVideoEncParams], self.ptr.data
+        ).type
+        return VideoEncParamsType(cython.cast(cython.int, t))
 
     @property
     def qp(self):
@@ -57,15 +66,17 @@ def qp(self):
         combined with `delta_qp` and the per-block delta in a manner
         documented for each type.
         """
-        return (<lib.AVVideoEncParams*> self.ptr.data).qp
-    
+        return cython.cast(cython.pointer[lib.AVVideoEncParams], self.ptr.data).qp
+
     @property
     def delta_qp(self):
         """
         Quantisation parameter offset from the base (per-frame) qp for a given
         plane (first index) and AC/DC coefficients (second index).
         """
-        cdef lib.AVVideoEncParams *p = <lib.AVVideoEncParams*> self.ptr.data
+        p: cython.pointer[lib.AVVideoEncParams] = cython.cast(
+            cython.pointer[lib.AVVideoEncParams], self.ptr.data
+        )
         return [[p.delta_qp[i][j] for j in range(2)] for i in range(4)]
 
     def block_params(self, idx):
@@ -85,25 +96,28 @@ def qp_map(self):
         """
         import numpy as np
 
-        cdef int mb_h = (self.frame.ptr.height + 15) // 16
-        cdef int mb_w = (self.frame.ptr.width + 15) // 16
-        cdef int nb_mb = mb_h * mb_w
-        cdef int block_idx
-        cdef int y
-        cdef int x
-        cdef VideoBlockParams block
+        mb_h: cython.int = (self.frame.ptr.height + 15) // 16
+        mb_w: cython.int = (self.frame.ptr.width + 15) // 16
+        nb_mb: cython.int = mb_h * mb_w
+        block_idx, x, y = cython.declare(cython.int)
+        block: VideoBlockParams
 
-        # Validate number of blocks
         if self.nb_blocks != nb_mb:
-            raise RuntimeError("Expected frame size to match number of blocks in side data")
-
-        # Validate type
-        cdef lib.AVVideoEncParamsType type = (<lib.AVVideoEncParams*> self.ptr.data).type
-        if type != lib.AVVideoEncParamsType.AV_VIDEO_ENC_PARAMS_MPEG2 and type != lib.AVVideoEncParamsType.AV_VIDEO_ENC_PARAMS_H264:
+            raise RuntimeError(
+                "Expected frame size to match number of blocks in side data"
+            )
+
+        type: lib.AVVideoEncParamsType = cython.cast(
+            cython.pointer[lib.AVVideoEncParams], self.ptr.data
+        ).type
+        if (
+            type != lib.AVVideoEncParamsType.AV_VIDEO_ENC_PARAMS_MPEG2
+            and type != lib.AVVideoEncParamsType.AV_VIDEO_ENC_PARAMS_H264
+        ):
             raise ValueError("Expected MPEG2 or H264")
 
         # Create a 2-D map with the number of macroblocks
-        cdef int32_t[:, ::1] map = np.empty((mb_h, mb_w), dtype=np.int32)
+        map = np.empty((mb_h, mb_w), dtype=np.int32)
 
         # Fill map with quantization parameter per macroblock
         for block_idx in range(nb_mb):
@@ -112,14 +126,19 @@ def qp_map(self):
             x = block.src_x // 16
             map[y, x] = self.qp + block.delta_qp
 
-        return np.asarray(map)
+        return map
 
 
-cdef class VideoBlockParams:
-    def __init__(self, VideoEncParams video_enc_params, int idx) -> None:
-        cdef uint8_t* base = <uint8_t*> video_enc_params.ptr.data
-        cdef Py_ssize_t offset = video_enc_params.blocks_offset + idx * video_enc_params.block_size
-        self.ptr = <lib.AVVideoBlockParams*> (base + offset)
+@cython.cclass
+class VideoBlockParams:
+    def __init__(self, video_enc_params: VideoEncParams, idx: cython.int) -> None:
+        base: cython.pointer[uint8_t] = cython.cast(
+            cython.pointer[uint8_t], video_enc_params.ptr.data
+        )
+        offset: cython.Py_ssize_t = (
+            video_enc_params.blocks_offset + idx * video_enc_params.block_size
+        )
+        self.ptr = cython.cast(cython.pointer[lib.AVVideoBlockParams], base + offset)
 
     def __repr__(self):
         return f"<av.sidedata.VideoBlockParams, src=({self.src_x}, {self.src_y}), size={self.w}x{self.h}, delta_qp={self.delta_qp}>"