From e24d3979fd326f7cbe20a2c6e6e817210f5e2ce7 Mon Sep 17 00:00:00 2001
From: Audionut <audionut11@gmail.com>
Date: Sun, 18 Aug 2024 14:45:09 +1000
Subject: [PATCH] Filter the pymediainfo output instead

---
 src/prep.py | 188 +++++++++++++++++++++++++++++++++++++++-------------
 1 file changed, 141 insertions(+), 47 deletions(-)

diff --git a/src/prep.py b/src/prep.py
index 5292b85a1..e52fc7d59 100644
--- a/src/prep.py
+++ b/src/prep.py
@@ -476,62 +476,156 @@ def get_video(self, videoloc, mode):
     Get and parse mediainfo
     """
     def exportInfo(self, video, isdir, folder_id, base_dir, export_text):
-        # Path to the MediaInfo CLI executable
-        mediainfo_cli_path = 'path/to/mediainfo/commandline.exe'  # Correct path to the MediaInfo CLI executable
+        def filter_mediainfo(data):
+            filtered = {
+                "creatingLibrary": data.get("creatingLibrary"),
+                "media": {
+                    "@ref": data["media"]["@ref"],
+                    "track": []
+                }
+            }
+            
+            for track in data["media"]["track"]:
+                if track["@type"] == "General":
+                    filtered["media"]["track"].append({
+                        "@type": track["@type"],
+                        "VideoCount": track.get("VideoCount"),
+                        "AudioCount": track.get("AudioCount"),
+                        "MenuCount": track.get("MenuCount"),
+                        "FileExtension": track.get("FileExtension"),
+                        "Format": track.get("Format"),
+                        "Format_Version": track.get("Format_Version"),
+                        "FileSize": track.get("FileSize"),
+                        "Duration": track.get("Duration"),
+                        "OverallBitRate": track.get("OverallBitRate"),
+                        "FrameRate": track.get("FrameRate"),
+                        "FrameCount": track.get("FrameCount"),
+                        "StreamSize": track.get("StreamSize"),
+                        "IsStreamable": track.get("IsStreamable"),
+                        "File_Created_Date": track.get("File_Created_Date"),
+                        "File_Created_Date_Local": track.get("File_Created_Date_Local"),
+                        "File_Modified_Date": track.get("File_Modified_Date"),
+                        "File_Modified_Date_Local": track.get("File_Modified_Date_Local"),
+                        "Encoded_Application": track.get("Encoded_Application"),
+                        "Encoded_Library": track.get("Encoded_Library"),
+                    })
+                elif track["@type"] == "Video":
+                    filtered["media"]["track"].append({
+                        "@type": track["@type"],
+                        "StreamOrder": track.get("StreamOrder"),
+                        "ID": track.get("ID"),
+                        "UniqueID": track.get("UniqueID"),
+                        "Format": track.get("Format"),
+                        "Format_Profile": track.get("Format_Profile"),
+                        "Format_Level": track.get("Format_Level"),
+                        "Format_Settings_CABAC": track.get("Format_Settings_CABAC"),
+                        "Format_Settings_RefFrames": track.get("Format_Settings_RefFrames"),
+                        "CodecID": track.get("CodecID"),
+                        "Duration": track.get("Duration"),
+                        "BitRate": track.get("BitRate"),
+                        "Width": track.get("Width"),
+                        "Height": track.get("Height"),
+                        "Sampled_Width": track.get("Sampled_Width"),
+                        "Sampled_Height": track.get("Sampled_Height"),
+                        "PixelAspectRatio": track.get("PixelAspectRatio"),
+                        "DisplayAspectRatio": track.get("DisplayAspectRatio"),
+                        "FrameRate_Mode": track.get("FrameRate_Mode"),
+                        "FrameRate": track.get("FrameRate"),
+                        "FrameCount": track.get("FrameCount"),
+                        "ColorSpace": track.get("ColorSpace"),
+                        "ChromaSubsampling": track.get("ChromaSubsampling"),
+                        "BitDepth": track.get("BitDepth"),
+                        "ScanType": track.get("ScanType"),
+                        "Delay": track.get("Delay"),
+                        "Delay_Source": track.get("Delay_Source"),
+                        "StreamSize": track.get("StreamSize"),
+                        "Encoded_Library": track.get("Encoded_Library"),
+                        "Encoded_Library_Name": track.get("Encoded_Library_Name"),
+                        "Encoded_Library_Version": track.get("Encoded_Library_Version"),
+                        "Encoded_Library_Settings": track.get("Encoded_Library_Settings"),
+                        "Language": track.get("Language"),
+                        "Default": track.get("Default"),
+                        "Forced": track.get("Forced"),
+                    })
+                elif track["@type"] == "Audio":
+                    filtered["media"]["track"].append({
+                        "@type": track["@type"],
+                        "StreamOrder": track.get("StreamOrder"),
+                        "ID": track.get("ID"),
+                        "UniqueID": track.get("UniqueID"),
+                        "Format": track.get("Format"),
+                        "Format_Settings_Mode": track.get("Format_Settings_Mode"),
+                        "Format_Settings_Endianness": track.get("Format_Settings_Endianness"),
+                        "CodecID": track.get("CodecID"),
+                        "Duration": track.get("Duration"),
+                        "BitRate_Mode": track.get("BitRate_Mode"),
+                        "BitRate": track.get("BitRate"),
+                        "Channels": track.get("Channels"),
+                        "ChannelPositions": track.get("ChannelPositions"),
+                        "ChannelLayout": track.get("ChannelLayout"),
+                        "SamplesPerFrame": track.get("SamplesPerFrame"),
+                        "SamplingRate": track.get("SamplingRate"),
+                        "SamplingCount": track.get("SamplingCount"),
+                        "FrameRate": track.get("FrameRate"),
+                        "BitDepth": track.get("BitDepth"),
+                        "Compression_Mode": track.get("Compression_Mode"),
+                        "Delay": track.get("Delay"),
+                        "Delay_Source": track.get("Delay_Source"),
+                        "Video_Delay": track.get("Video_Delay"),
+                        "StreamSize": track.get("StreamSize"),
+                        "Language": track.get("Language"),
+                        "Default": track.get("Default"),
+                        "Forced": track.get("Forced"),
+                    })
+                elif track["@type"] == "Text":
+                    filtered["media"]["track"].append({
+                        "@type": track["@type"],
+                        "@typeorder": track.get("@typeorder"),
+                        "StreamOrder": track.get("StreamOrder"),
+                        "ID": track.get("ID"),
+                        "UniqueID": track.get("UniqueID"),
+                        "Format": track.get("Format"),
+                        "CodecID": track.get("CodecID"),
+                        "Duration": track.get("Duration"),
+                        "BitRate": track.get("BitRate"),
+                        "FrameRate": track.get("FrameRate"),
+                        "FrameCount": track.get("FrameCount"),
+                        "ElementCount": track.get("ElementCount"),
+                        "StreamSize": track.get("StreamSize"),
+                        "Title": track.get("Title"),
+                        "Language": track.get("Language"),
+                        "Default": track.get("Default"),
+                        "Forced": track.get("Forced"),
+                    })
+                elif track["@type"] == "Menu":
+                    filtered["media"]["track"].append({
+                        "@type": track["@type"],
+                        "extra": track.get("extra"),
+                    })
+            
+            return filtered
 
         if not os.path.exists(f"{base_dir}/tmp/{folder_id}/MEDIAINFO.txt") and export_text:
             console.print("[bold yellow]Exporting MediaInfo...")
-            
-            # MediaInfo to text
             if not isdir:
                 os.chdir(os.path.dirname(video))
-            
-            # Use MediaInfo CLI to get text output with utf-8 encoding
-            result = subprocess.run(
-                [mediainfo_cli_path, '--Output=TEXT', video],
-                stdout=subprocess.PIPE,
-                stderr=subprocess.PIPE,
-                encoding='utf-8'
-            )
-            
-            if result.returncode != 0:
-                console.print(f"[bold red]Error: {result.stderr}")
-                return None
-
-            media_info = result.stdout
-            
-            if media_info:
-                with open(f"{base_dir}/tmp/{folder_id}/MEDIAINFO.txt", 'w', newline="", encoding='utf-8') as export:
-                    export.write(media_info)
-
-                # Save the MediaInfo with clean path
-                with open(f"{base_dir}/tmp/{folder_id}/MEDIAINFO_CLEANPATH.txt", 'w', newline="", encoding='utf-8') as export_cleanpath:
-                    export_cleanpath.write(media_info.replace(video, os.path.basename(video)))
-
+            media_info = MediaInfo.parse(video, output="STRING", full=False, mediainfo_options={'inform_version': '1'})
+            with open(f"{base_dir}/tmp/{folder_id}/MEDIAINFO.txt", 'w', newline="", encoding='utf-8') as export:
+                export.write(media_info)
+            with open(f"{base_dir}/tmp/{folder_id}/MEDIAINFO_CLEANPATH.txt", 'w', newline="", encoding='utf-8') as export_cleanpath:
+                export_cleanpath.write(media_info.replace(video, os.path.basename(video)))
             console.print("[bold green]MediaInfo Exported.")
 
-        if not os.path.exists(f"{base_dir}/tmp/{folder_id}/MediaInfo.json"):
-            # MediaInfo to JSON using the CLI
-            result = subprocess.run(
-                [mediainfo_cli_path, '--Output=JSON', video],
-                stdout=subprocess.PIPE,
-                stderr=subprocess.PIPE,
-                encoding='utf-8'
-            )
-            
-            if result.returncode != 0:
-                console.print(f"[bold red]Error: {result.stderr}")
-                return None
-            
-            media_info_json = result.stdout
-            
-            if media_info_json:
-                with open(f"{base_dir}/tmp/{folder_id}/MediaInfo.json", 'w', encoding='utf-8') as export:
-                    export.write(media_info_json)
-            
+        if not os.path.exists(f"{base_dir}/tmp/{folder_id}/MediaInfo.json.txt"):
+            media_info_json = MediaInfo.parse(video, output="JSON", mediainfo_options={'inform_version': '1'})
+            media_info_dict = json.loads(media_info_json)
+            filtered_info = filter_mediainfo(media_info_dict)
+            with open(f"{base_dir}/tmp/{folder_id}/MediaInfo.json", 'w', encoding='utf-8') as export:
+                json.dump(filtered_info, export, indent=4)
+
         with open(f"{base_dir}/tmp/{folder_id}/MediaInfo.json", 'r', encoding='utf-8') as f:
             mi = json.load(f)
-
+        
         return mi