From f68c38a4803da61fda1344f01260ea49f1577dce 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  | 14 +++++++++-----
 av/video/stream.pyi   |  7 ++++---
 tests/test_streams.py | 10 ++++++++++
 5 files changed, 26 insertions(+), 9 deletions(-)

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..95ec1b551 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,13 @@ 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
+        elif isinstance(value, ThreadType):
+            self.ptr.thread_type = value.value
+        raise TypeError(f"Invalid thread_type setter: {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/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)