Fix type stubs in subtitle.pyi

PyAV-Org · Mar 4, 2024 · ef97595 · ef97595
1 parent 1a1a6e2
commit ef97595
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 8 deletions.
diff --git a/av/subtitles/subtitle.pyi b/av/subtitles/subtitle.pyi
@@ -1,6 +1,4 @@
-from typing import Any, Iterator, Literal
-
-from .stream import SubtitleStream
+from typing import Iterator, Literal
 
 class SubtitleSet:
     format: int
@@ -13,15 +11,15 @@ class SubtitleSet:
     def __getitem__(self, i: int) -> Subtitle: ...
 
 class Subtitle:
-    type: Literal["none", "bitmap", "text", "ass"]
+    type: Literal[b"none", b"bitmap", b"text", b"ass"]
 
 class BitmapSubtitle(Subtitle):
-    type: Literal["bitmap"]
+    type: Literal[b"bitmap"]
     x: int
     y: int
     width: int
     height: int
-    nb_colors: Any
+    nb_colors: int
     planes: tuple[BitmapSubtitlePlane, ...]
 
 class BitmapSubtitlePlane:
@@ -30,9 +28,9 @@ class BitmapSubtitlePlane:
     buffer_size: int
 
 class TextSubtitle(Subtitle):
-    type: Literal["text"]
+    type: Literal[b"text"]
     text: str
 
 class AssSubtitle(Subtitle):
-    type: Literal["text"]
+    type: Literal[b"ass"]
     ass: str
diff --git a/tests/test_subtitles.py b/tests/test_subtitles.py
@@ -23,6 +23,7 @@ def test_movtext(self):
 
         sub = subset[0]
         self.assertIsInstance(sub, AssSubtitle)
+        self.assertEqual(sub.type, b"ass")
         self.assertEqual(sub.ass, "0,0,Default,,0,0,0,,- Test 1.\\N- Test 2.")
 
     def test_vobsub(self):
@@ -43,10 +44,12 @@ def test_vobsub(self):
 
         sub = subset[0]
         self.assertIsInstance(sub, BitmapSubtitle)
+        self.assertEqual(sub.type, b"bitmap")
         self.assertEqual(sub.x, 259)
         self.assertEqual(sub.y, 379)
         self.assertEqual(sub.width, 200)
         self.assertEqual(sub.height, 24)
+        self.assertEqual(sub.nb_colors, 4)
 
         bms = sub.planes
         self.assertEqual(len(bms), 1)