From 5fec2c936d3cde6d77d2b4f4509a89c66f3b0025 Mon Sep 17 00:00:00 2001
From: Scott Alfter <scott@alfter.us>
Date: Mon, 18 Nov 2024 13:40:25 -0800
Subject: [PATCH] improve subtitle selection logic: can prefer hearing-impaired
 track or regular track

---
 README.md      |  6 ++++--
 pickstreams.py | 54 ++++++++++++++++++++++++++++++++++++++++++++++----
 2 files changed, 54 insertions(+), 6 deletions(-)

diff --git a/README.md b/README.md
index a7dd487..031fce8 100644
--- a/README.md
+++ b/README.md
@@ -8,7 +8,9 @@ a frontend to ffprobe that:
 
 Usage:
 
-pickstreams.py <file> <audio_lang> <sub_lang>
+pickstreams.py <file> <audio_lang> <sub_lang> <hi_preferred>
 
 If audio_lang is left empty, all audio strrams are shown. If sub_lang is
-left empty, all subtitle streams are shown.
+left empty, all subtitle streams are shown.  If hi_preferred is set to "hi",
+hearing-impaired subtitle tracks will be preferred; if left empty, regular
+subtitle tracks will be preferred.
diff --git a/pickstreams.py b/pickstreams.py
index 9983eea..505a479 100755
--- a/pickstreams.py
+++ b/pickstreams.py
@@ -6,6 +6,24 @@ import json
 
 from pprint import pprint
 
+try:
+  audio_lang=sys.argv[2]
+except IndexError:
+  audio_lang=""
+try:
+  sub_lang=sys.argv[3]
+except IndexError:
+  sub_lang=""
+try:
+  if sys.argv[4]=="hi":
+    prefer_hi=True
+  else:
+    prefer_hi=False
+except IndexError:
+  prefer_hi=False
+audio_info=""
+sub_info=""
+sub_hi_info=""
 streams=json.loads(subprocess.run(["ffprobe", "-v", "error", "-show_streams", "-print_format", "json", sys.argv[1]], capture_output=True).stdout.decode("utf-8"))["streams"]
 for i in streams:
   matched=True
@@ -19,9 +37,10 @@ for i in streams:
           out=out+" "+str(int(i["tags"]["BPS"])/1000)
         except KeyError:
           out=out+" unknown"
+      print(out)
     case "audio":
       try:
-        if sys.argv[2]!="" and sys.argv[2]!=i["tags"]["language"]:
+        if audio_lang!="" and audio_lang!=i["tags"]["language"]:
           matched=False
       except:
         pass
@@ -36,9 +55,14 @@ for i in streams:
       out=out+" "+i["tags"]["language"]
       if i["disposition"]["default"]==1:
         out=out+" default"
+      if matched==True:
+        audio_info=out
+      if audio_lang=="":
+        print(out)
     case "subtitle":
+      hi_flag=False
       try:
-        if sys.argv[3]!="" and sys.argv[3]!=i["tags"]["language"]:
+        if sub_lang!="" and sub_lang!=i["tags"]["language"]:
           matched=False
       except:
         pass
@@ -47,12 +71,34 @@ for i in streams:
         out=out+" default"
       if i["disposition"]["hearing_impaired"]==1:
         out=out+" hi"
+        if matched==True:
+          sub_hi_info=out
+        hi_flag=True
       try:
         if i["tags"]["title"].find("SDH")!=-1:
           out=out+" hi"
+          if matched==True:
+            sub_hi_info=out
+            hi_flag=True
       except KeyError:
         pass
+      if hi_flag==False and matched==True:
+        sub_info=out
+      if sub_lang=="":
+        print(out)
     case _:
       out=str(i["index"])+" "+i["codec_type"]+" "+i["codec_name"]
-  if matched==True:
-    print(out)    
+if audio_info!="" and audio_lang!="":
+  print(audio_info)
+if prefer_hi==False:
+  if sub_info!="" and sub_lang!="":
+    print(sub_info)
+  else:
+    if sub_hi_info!="" and sub_lang!="":
+      print(sub_hi_info)
+else:
+  if sub_hi_info!="" and sub_lang!="":
+    print(sub_hi_info)
+  else:
+    if sub_info!="" and sub_lang!="":
+      print(sub_info)