Allow extradata to be set for encoders

PyAV-Org · Sep 12, 2024 · ee5c40e · ee5c40e
1 parent 8809032
commit ee5c40e
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 12 deletions.
diff --git a/av/codec/context.pyx b/av/codec/context.pyx
@@ -189,9 +189,6 @@ cdef class CodecContext:
 
     @extradata.setter
     def extradata(self, data):
-        if not self.is_decoder:
-            raise ValueError("Can only set extradata for decoders.")
-
         if data is None:
             lib.av_freep(&self.ptr.extradata)
             self.ptr.extradata_size = 0

diff --git a/av/container/output.pyx b/av/container/output.pyx
@@ -93,10 +93,14 @@ cdef class OutputContainer(Container):
         # Now lets set some more sane video defaults
         elif codec.type == lib.AVMEDIA_TYPE_VIDEO:
             codec_context.pix_fmt = lib.AV_PIX_FMT_YUV420P
-            codec_context.width = 640
-            codec_context.height = 480
-            codec_context.bit_rate = 1024000
-            codec_context.bit_rate_tolerance = 128000
+            codec_context.width = kwargs.pop("width", 640)
+            codec_context.height = kwargs.pop("height", 480)
+            codec_context.bit_rate = kwargs.pop("bit_rate", 1024000)
+            codec_context.bit_rate_tolerance = kwargs.pop("bit_rate_tolerance", 128000)
+            try:
+                to_avrational(kwargs.pop("time_base"), &codec_context.time_base)
+            except KeyError:
+                pass
             to_avrational(rate or 24, &codec_context.framerate)
 
             stream.avg_frame_rate = codec_context.framerate
@@ -105,9 +109,14 @@ cdef class OutputContainer(Container):
         # Some sane audio defaults
         elif codec.type == lib.AVMEDIA_TYPE_AUDIO:
             codec_context.sample_fmt = codec.sample_fmts[0]
-            codec_context.bit_rate = 128000
-            codec_context.bit_rate_tolerance = 32000
+            codec_context.bit_rate = kwargs.pop("bit_rate", 128000)
+            codec_context.bit_rate_tolerance = kwargs.pop("bit_rate_tolerance", 32000)
+            try:
+                to_avrational(kwargs.pop("time_base"), &codec_context.time_base)
+            except KeyError:
+                pass
             codec_context.sample_rate = rate or 48000
+            stream.time_base = codec_context.time_base
             lib.av_channel_layout_default(&codec_context.ch_layout, 2)
 
         # Some formats want stream headers to be separate

diff --git a/tests/test_codec_context.py b/tests/test_codec_context.py
@@ -105,9 +105,9 @@ def test_encoder_extradata(self):
         self.assertEqual(ctx.extradata, None)
         self.assertEqual(ctx.extradata_size, 0)
 
-        with self.assertRaises(ValueError) as cm:
-            ctx.extradata = b"123"
-        self.assertEqual(str(cm.exception), "Can only set extradata for decoders.")
+        ctx.extradata = b"123"
+        self.assertEqual(ctx.extradata, b"123")
+        self.assertEqual(ctx.extradata_size, 3)
 
     def test_encoder_pix_fmt(self):
         ctx = av.codec.Codec("h264", "w").create()