Optimizations to stream classes

src/models/streams.py

@@ -1,15 +1,22 @@
 from pydantic import BaseModel, Field
-from typing import Optional
+from enum import Enum
 from datetime import datetime
 from pathlib import Path
 from functools import cached_property
 
 from lib.fs import sanitize_filename
 
+class OriginalMediaType(str, Enum):
+    DVD = "DVD"
+    BluRay = "BluRay"
+    TVRip = "TVRip"
+    WebRip = "WebRip"
+    Other = "Other"
+
 
 class Tags(BaseModel):
-    title: Optional[str] = Field(None)
+    title: str | None = Field(None)
-    language: Optional[str] = Field(None)
+    language: str | None = Field(None)
 
 
 class Disposition(BaseModel):
@@ -18,24 +25,31 @@ class Disposition(BaseModel):
 
 
 class VideoStream(BaseModel):
-    index: int = Field(None)
+    index: int | None = Field(None)
-    width: int = Field(None)
+    width: int | None = Field(None)
-    height: int = Field(None)
+    height: int | None = Field(None)
-    avg_frame_rate: str = Field(None, alias="avg_frame_rate")
+    avg_frame_rate: str | None = Field(None, alias="avg_frame_rate")
-    codec_name: str = Field(None)
+    codec_name: str | None = Field(None)
     disposition: Disposition
     tags: Tags | None = Field(None)
     preset: int = Field(2)
     crf: int = Field(30)
-    original_media_type: str = Field(None)
+    original_media_type: str | None = Field(None)
     film_grain: int = Field(20)
 
     @cached_property
     def fps(self) -> float:
-        frame_numbers = [int(x) for x in self.avg_frame_rate.split("/")]
+        """Calculate frames per second from avg_frame_rate string."""
-        return frame_numbers[0] / frame_numbers[1]
+        try:
+            if not self.avg_frame_rate:
+                return 0.0
+            num, denom = (int(x) for x in self.avg_frame_rate.split("/"))
+            return num / denom if denom else 0.0
+        except Exception:
+            return 0.0
 
     def cmd_flags(self) -> list[str]:
+        """Generate ffmpeg command flags for this video stream."""
         return [
             "-map",
             f"0:{self.index}",
@@ -50,7 +64,7 @@ class VideoStream(BaseModel):
             "-vf",
             "'scale=iw*sar:ih,setsar=1,scale=-2:ih:lanczos'",
             "-g",
-            f"{round(self.fps) * 5}",
+            str(round(self.fps) * 5),
         ]
 
 
@@ -63,20 +77,24 @@ class AudioStream(BaseModel):
     tags: Tags | None = Field(None)
 
     def cmd_flags(self, index: int) -> list[str]:
+        """Generate ffmpeg command flags for this audio stream."""
         side_loaded = []
+        if "side" in self.channel_layout:
+            side_loaded = [f"-filter:a:{index}", "'channelmap=channel_layout=5.1'"]
         default_disposition = (
             [f"-disposition:a:{index}", "default"] if self.disposition.default else []
         )
-        if "side" in self.channel_layout:
+        language = getattr(self.tags, 'language', None) if self.tags else None
-            side_loaded = [f"-filter:a:{index}", "'channelmap=channel_layout=5.1'"]
+        title = getattr(self.tags, 'title', None) if self.tags else None
+        meta_flags = []
+        if language:
+            meta_flags += [f"-metadata:s:a:{index}", f"language={language}"]
+        if title:
+            meta_flags += [f"-metadata:s:a:{index}", f"title='{title}'"]
         return (
             ["-map", f"0:{self.index}", f"-c:a:{index}", "libopus"]
             + side_loaded
-            + [f"-metadata:s:a:{index}", f"language={self.tags.language}"]
+            + meta_flags
-            + [
-                f"-metadata:s:a:{index}",
-                f"title='{self.tags.title}'",
-            ]
             + default_disposition
         )
 
@@ -88,23 +106,27 @@ class SubtitleStream(BaseModel):
     tags: Tags | None = Field(None)
 
     def cmd_flags(self, index: int) -> list[str]:
-        return ["-map", f"0:s:{index}", f"-c:s:{index}", "copy"] + [
+        """Generate ffmpeg command flags for this subtitle stream."""
-            f"-metadata:s:s:{index}",
+        language = getattr(self.tags, 'language', None) if self.tags else None
-            f"language={self.tags.language}",
+        meta_flags = []
-        ]
+        if language:
+            meta_flags += [f"-metadata:s:s:{index}", f"language={language}"]
+        return ["-map", f"0:s:{index}", f"-c:s:{index}", "copy"] + meta_flags
 
 
 class VideoFile(BaseModel):
     video_path: Path
-    title: str = Field(None)
+    title: str | None = Field(None)
-    date_released: Optional[datetime] = Field(None)
+    date_released: datetime | None = Field(None)
-    imdb: Optional[str] = Field(None)
+    imdb: str | None = Field(None)
-    video_streams: list[VideoStream] = Field([])
+    video_streams: list[VideoStream] = Field(default_factory=list)
-    audio_streams: list[AudioStream] = Field([])
+    audio_streams: list[AudioStream] = Field(default_factory=list)
-    subtitle_streams: list[SubtitleStream] = Field([])
+    subtitle_streams: list[SubtitleStream] = Field(default_factory=list)
     allow_subtitle: bool = Field(True)
+    original_media_type: OriginalMediaType = Field(OriginalMediaType.DVD)
 
     def cmd(self) -> list[str]:
+        """Generate the full ffmpeg command for this video file."""
         cmd = [
             "ffmpeg",
             "-hide_banner",
@@ -114,17 +136,15 @@ class VideoFile(BaseModel):
             f"'{str(self.video_path.absolute())}'",
             "-map_metadata",
             "-1",
-            "-metadata",
-            f"title='{self.title}'",
-            "-metadata",
-            f"date_released={self.date_released.strftime('%Y-%m-%d')}",
-            "-metadata",
-            f"imdb={self.imdb}",
-            "-metadata",
-            "original_media_type=DVD",
-            "-pix_fmt",
-            "yuv420p10le",
         ]
+        if self.title:
+            cmd += ["-metadata", f"title='{self.title}'"]
+        if self.date_released:
+            cmd += ["-metadata", f"date_released={self.date_released.strftime('%Y-%m-%d')}"]
+        if self.imdb:
+            cmd += ["-metadata", f"imdb={self.imdb}"]
+        cmd += ["-metadata", f"original_media_type={self.original_media_type.value}"]
+        cmd += ["-pix_fmt", "yuv420p10le"]
         for stream in self.video_streams:
             cmd.extend(stream.cmd_flags())
         for i, stream in enumerate(self.audio_streams):
@@ -132,12 +152,10 @@ class VideoFile(BaseModel):
         if self.allow_subtitle:
             for i, stream in enumerate(self.subtitle_streams):
                 cmd.extend(stream.cmd_flags(i))
-
+        year = self.date_released.year if self.date_released else "unknown"
         output_file = (
             self.video_path.absolute().parent
-            / f"{sanitize_filename(self.title)}.{self.date_released.year}.mkv"
+            / f"{sanitize_filename(self.title or 'untitled')}.{year}.mkv"
         )
-
+        cmd.append(f"'{str(output_file)}'")
-        cmd.extend([f"'{str(output_file)}'"])
-
         return cmd