From 906c2e9cb9570492a021d20029fe4facf201ff05 Mon Sep 17 00:00:00 2001
From: Vitiko <averroista@protonmail.com>
Date: Thu, 14 Sep 2023 15:58:18 -0400
Subject: [PATCH] EmbeddedSubtitles provider: improve cache management (Fix
 #2241)

---
 libs/fese/__init__.py                                |  2 +-
 libs/fese/container.py                               | 12 ++++++++++--
 libs/subliminal_patch/providers/embeddedsubtitles.py |  7 +++++++
 3 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/libs/fese/__init__.py b/libs/fese/__init__.py
index 132e10ff3..4b0020107 100755
--- a/libs/fese/__init__.py
+++ b/libs/fese/__init__.py
@@ -4,4 +4,4 @@
 from .container import FFprobeVideoContainer
 from .stream import FFprobeSubtitleStream
 
-__version__ = "0.2.8"
+__version__ = "0.2.9"
diff --git a/libs/fese/container.py b/libs/fese/container.py
index d65aee17f..57bd7fe5e 100755
--- a/libs/fese/container.py
+++ b/libs/fese/container.py
@@ -81,6 +81,7 @@ class FFprobeVideoContainer:
         overwrite=True,
         timeout=600,
         convert_format=None,
+        basename_callback=None,
     ):
         """Extracts a list of subtitles converting them. Returns a dictionary of the
         extracted filenames by index.
@@ -95,6 +96,8 @@ class FFprobeVideoContainer:
         :param timeout: subprocess timeout in seconds (default: 600)
         :param convert_format: format to convert selected subtitles. Defaults to
         srt
+        :param basename_callback: a callback that takes the filename path. Only used if
+        custom_dir is set. Defaults to `os.path.basename`
         :raises: ExtractionError, UnsupportedCodec, OSError
         """
         extract_command = [FFMPEG_PATH, "-v", FF_LOG_LEVEL]
@@ -116,7 +119,8 @@ class FFprobeVideoContainer:
                 f"{os.path.splitext(self.path)[0]}.{subtitle.suffix}.{extension_to_use}"
             )
             if custom_dir is not None:
-                sub_path = os.path.join(custom_dir, os.path.basename(sub_path))
+                basename_callback = basename_callback or os.path.basename
+                sub_path = os.path.join(custom_dir, basename_callback(sub_path))
 
             if not overwrite and sub_path in collected_paths:
                 sub_path = f"{os.path.splitext(sub_path)[0]}.{len(collected_paths):02}.{extension_to_use}"
@@ -156,6 +160,7 @@ class FFprobeVideoContainer:
         overwrite=True,
         timeout=600,
         fallback_to_convert=True,
+        basename_callback=None,
     ):
         """Extracts a list of subtitles with ffmpeg's copy method. Returns a dictionary
         of the extracted filenames by index.
@@ -167,6 +172,8 @@ class FFprobeVideoContainer:
         :param timeout: subprocess timeout in seconds (default: 600)
         :param fallback_to_convert: fallback to stream's default convert format if it is
         incompatible with copy
+        :param basename_callback: a callback that takes the filename path. Only used if
+        custom_dir is set. Defaults to `os.path.basename`
         :raises: ExtractionError, UnsupportedCodec, OSError
         """
         extract_command = [FFMPEG_PATH, "-v", FF_LOG_LEVEL]
@@ -184,7 +191,8 @@ class FFprobeVideoContainer:
         for subtitle in subtitles:
             sub_path = f"{os.path.splitext(self.path)[0]}.{subtitle.suffix}.{subtitle.extension}"
             if custom_dir is not None:
-                sub_path = os.path.join(custom_dir, os.path.basename(sub_path))
+                basename_callback = basename_callback or os.path.basename
+                sub_path = os.path.join(custom_dir, basename_callback(sub_path))
 
             if not overwrite and sub_path in collected_paths:
                 sub_path = f"{os.path.splitext(sub_path)[0]}.{len(collected_paths):02}.{subtitle.extension}"
diff --git a/libs/subliminal_patch/providers/embeddedsubtitles.py b/libs/subliminal_patch/providers/embeddedsubtitles.py
index a5c49df91..76e8afdf4 100644
--- a/libs/subliminal_patch/providers/embeddedsubtitles.py
+++ b/libs/subliminal_patch/providers/embeddedsubtitles.py
@@ -2,6 +2,7 @@
 
 import functools
 import logging
+import hashlib
 import os
 import re
 import shutil
@@ -214,6 +215,7 @@ class EmbeddedSubtitlesProvider(Provider):
                 self._cache_dir,
                 timeout=self._timeout,
                 fallback_to_convert=True,
+                basename_callback=_basename_callback,
             )
             # Add the extracted paths to the containter path key
             self._cached_paths[container.path] = extracted
@@ -345,6 +347,11 @@ def _get_pretty_release_name(stream, container):
     return f"{os.path.splitext(bname)[0]}.{stream.suffix}"
 
 
+def _basename_callback(path: str):
+    path, ext = os.path.splitext(path)
+    return hashlib.md5(path.encode()).hexdigest() + ext
+
+
 # TODO: improve this
 _SIGNS_LINE_RE = re.compile(r",([\w|_]{,15}(sign|fx|karaoke))", flags=re.IGNORECASE)