From 9f51a441847b559789c5dcf4ac371815932437ab Mon Sep 17 00:00:00 2001
From: WyattBlue <wyattblue@auto-editor.com>
Date: Mon, 18 Nov 2024 13:33:47 -0500
Subject: [PATCH] Fix thread_type setter

---
 av/about.py           |   2 +-
 av/codec/context.pyi  |   2 ++
 av/codec/context.pyx  |  13 ++++++++-----
 av/video/stream.pyi   |   7 ++++---
 data.ts               | Bin 0 -> 1128 bytes
 tests/test_streams.py |  10 ++++++++++
 6 files changed, 25 insertions(+), 9 deletions(-)
 create mode 100644 data.ts

diff --git a/av/about.py b/av/about.py
index a3d32212b..73437bbc8 100644
--- a/av/about.py
+++ b/av/about.py
@@ -1 +1 @@
-__version__ = "14.0.0rc1"
+__version__ = "14.0.0rc2"
diff --git a/av/codec/context.pyi b/av/codec/context.pyi
index 8600deaaa..8b7458597 100644
--- a/av/codec/context.pyi
+++ b/av/codec/context.pyi
@@ -12,6 +12,8 @@ class ThreadType(Flag):
     FRAME: ClassVar[ThreadType]
     SLICE: ClassVar[ThreadType]
     AUTO: ClassVar[ThreadType]
+    def __get__(self, i: object | None, owner: type | None = None) -> ThreadType: ...
+    def __set__(self, instance: object, value: int | str | ThreadType) -> None: ...
 
 class SkipType(Enum):
     NONE: ClassVar[SkipType]
diff --git a/av/codec/context.pyx b/av/codec/context.pyx
index 528a73872..6e039a4e1 100644
--- a/av/codec/context.pyx
+++ b/av/codec/context.pyx
@@ -138,10 +138,8 @@ cdef class CodecContext:
         self.codec = wrap_codec(codec if codec != NULL else self.ptr.codec)
 
         # Set reasonable threading defaults.
-        # count == 0 -> use as many threads as there are CPUs.
-        # type == 2 -> thread within a frame. This does not change the API.
-        self.ptr.thread_count = 0
-        self.ptr.thread_type = 2
+        self.ptr.thread_count = 0  # use as many threads as there are CPUs.
+        self.ptr.thread_type = 0x02  # thread within a frame. Does not change the API.
 
     def _get_flags(self):
         return self.ptr.flags
@@ -623,7 +621,12 @@ cdef class CodecContext:
     def thread_type(self, value):
         if self.is_open:
             raise RuntimeError("Cannot change thread_type after codec is open.")
-        self.ptr.thread_type = value.value
+        if type(value) is int:
+            self.ptr.thread_type = value
+        elif type(value) is str:
+            self.ptr.thread_type = ThreadType[value].value
+        else:
+            self.ptr.thread_type = value.value
 
     @property
     def skip_frame(self):
diff --git a/av/video/stream.pyi b/av/video/stream.pyi
index f0cdd5eb4..dd670d3cf 100644
--- a/av/video/stream.pyi
+++ b/av/video/stream.pyi
@@ -1,6 +1,7 @@
 from fractions import Fraction
-from typing import Any, Iterator, Literal
+from typing import Iterator, Literal
 
+from av.codec.context import ThreadType
 from av.packet import Packet
 from av.stream import Stream
 
@@ -12,8 +13,6 @@ class VideoStream(Stream):
     bit_rate: int | None
     max_bit_rate: int | None
     bit_rate_tolerance: int
-    thread_count: int
-    thread_type: Any
     sample_aspect_ratio: Fraction | None
     display_aspect_ratio: Fraction | None
     codec_context: VideoCodecContext
@@ -24,6 +23,8 @@ class VideoStream(Stream):
 
     # from codec context
     format: VideoFormat
+    thread_count: int
+    thread_type: ThreadType
     width: int
     height: int
     bits_per_coded_sample: int
diff --git a/data.ts b/data.ts
new file mode 100644
index 0000000000000000000000000000000000000000..a33b8b31531f8b173bfb0b5d0098a72173fe9d19
GIT binary patch
literal 1128
zcmZ>F5ENi=`k>0dc#whNKjVJ}#y<@r9zu+4Zf>~+sp*`-sYPX($*Bg0<u%Ty{tqtT
z?!X|xz_5WA>~;pm4-8rxSvCzWC*pLCdjQbwOdEv2Zhy$|fq@M~?5kf8gj36)k#Kip
zFqj7nc?O0$hC$;fREIDC(=mf^O9QJS10w@tNosM4LP}yuqJkmN?S>Gy{~Ho+XGXZ)
T2<Ubru-h4jM!2&g+-?j2(Q?e4

literal 0
HcmV?d00001

diff --git a/tests/test_streams.py b/tests/test_streams.py
index b1f2bd12f..37d1816cd 100644
--- a/tests/test_streams.py
+++ b/tests/test_streams.py
@@ -21,6 +21,16 @@ def test_selection(self) -> None:
             fate_suite("amv/MTV_high_res_320x240_sample_Penguin_Joke_MTV_from_WMV.amv")
         )
         video = container.streams.video[0]
+
+        video.thread_type = av.codec.context.ThreadType.AUTO
+        assert video.thread_type == av.codec.context.ThreadType.AUTO
+
+        video.thread_type = 0x03
+        assert video.thread_type == av.codec.context.ThreadType.AUTO
+
+        video.thread_type = "AUTO"
+        assert video.thread_type == av.codec.context.ThreadType.AUTO
+
         audio = container.streams.audio[0]
 
         assert [video] == container.streams.get(video=0)