plexflow 0.0.64__py3-none-any.whl

plexflow/utils/subtitle/search.py

@@ -0,0 +1,76 @@
+import re
+
+class SubtitleSearcher:
+    """
+    A class used to search for subtitles in a text.
+
+    ...
+
+    Attributes
+    ----------
+    hint_words : set
+        a set of words to be used as hints in the subtitle search
+
+    Methods
+    -------
+    compile_subtitle_hint_pattern(hint_word):
+        Returns a compiled regular expression pattern based on the provided hint word.
+    search_subtitles(text):
+        Searches for subtitles in the provided text and returns a list of matches.
+    """
+
+    def __init__(self, hint_words):
+        """
+        Constructs all the necessary attributes for the SubtitleSearcher object.
+
+        Parameters
+        ----------
+            hint_words : set
+                a set of words to be used as hints in the subtitle search
+        """
+        self.hint_words = hint_words
+
+    def compile_subtitle_hint_pattern(self, hint_word):
+        """
+        Returns a compiled regular expression pattern based on the provided hint word.
+
+        The pattern searches for the hint word preceded by 'text' or 'subtitles' and 
+        followed by any characters. The search is case-insensitive and includes newline 
+        characters.
+
+        Parameters
+        ----------
+            hint_word : str
+                a word to be used as a hint in the subtitle search
+
+        Returns
+        -------
+            pattern : re.Pattern
+                a compiled regular expression pattern
+        """
+        pattern = re.compile(rf'\b(text|subtitles?)\b(.*?)\b({hint_word})\b', re.IGNORECASE | re.DOTALL | re.UNICODE)
+        return pattern
+
+    def search_subtitles(self, text):
+        """
+        Searches for subtitles in the provided text and returns a list of matches.
+
+        The search is performed using the hint words provided during the object 
+        initialization. Each hint word is used to compile a regular expression pattern, 
+        which is then used to search the text.
+
+        Parameters
+        ----------
+            text : str
+                the text to be searched
+
+        Returns
+        -------
+            matches : list
+                a list of matches found in the text
+        """
+        matches = []
+        for hint_word in self.hint_words:
+            pattern = self.compile_subtitle_hint_pattern(hint_word)
+            matches.extend([m.group(3) for m in re.finditer(pattern, text)])
+        return ",".join(set([m.strip().lower() for m in matches]))

plexflow/utils/tasks/decorators.py

@@ -0,0 +1,78 @@
+import inspect
+from plexflow.core.context.partial_context import PartialContext
+import logging
+import os
+import json
+
+def plexflow(task_func):
+    def wrapper(*args, **kwargs):
+        task_mode = os.getenv("TASK_MODE", "default")
+        default_ttl = 3600
+
+        if task_mode == "k8s":
+            dag_run_id = os.getenv("AIRFLOW_RUN_ID", None)
+            # setup logging level
+            logging.basicConfig(level=logging.INFO)
+            
+            # Loop over environment variables and look for
+            # variables that start with ARG_ or KW_ARG_
+            # These are the extra arguments and keyword arguments
+            # that are passed to the task
+            # The ARG_ variables are passed as positional arguments
+            # and must be sorted from ARG_0 to ARG_N
+            # The KW_ARG_ variables are passed as keyword arguments
+            # and are ordered arbitrarily
+            extra_args = []
+            extra_kwargs = {}
+            # lets first sort the env variables by the name of the variable
+            sorted_env = sorted(os.environ.items(), key=lambda x: x[0])
+            logging.info(f"Sorted env: {sorted_env}")
+            
+            for key, value in sorted_env:
+                if key.startswith("ARG_"):
+                    extra_args.append(value)
+                elif key.startswith("KW_ARG_"):
+                    arg = json.loads(value)
+                    extra_kwargs[arg["key"]] = arg["value"]
+            
+            logging.info(f"Extra args: {extra_args}")
+            logging.info(f"Extra kwargs: {extra_kwargs}")
+            
+            # Now lets update the args and kwargs with the extra arguments
+            args = list(args) + extra_args
+            kwargs.update(extra_kwargs)
+        else:
+            context = kwargs.get('ti', None)
+            dag_run_id = context.run_id
+
+        context_id = PartialContext.create_universal_id(dag_run_id)
+
+        logging.info(f"Task mode: {task_mode}")
+        logging.info(f"Context id: {context_id}")
+        logging.info(f"Dag run id: {dag_run_id}")
+        logging.info(f"Default TTL: {default_ttl}")
+        logging.info(f"Universal id: {context_id}") 
+        
+        sig = inspect.signature(task_func)
+        logging.info(f"Function signature: {sig.parameters}")
+        
+        func_kwargs = {}
+        pos_arg_index = 0
+
+        for param_name, param in sig.parameters.items():
+            arg_type = param.annotation
+            
+            if param_name in kwargs:
+                func_kwargs[param_name] = kwargs[param_name]
+            elif issubclass(arg_type, PartialContext):
+                # check if arg_type is subclass of PartialContext
+                # Create an instance of the class
+                func_kwargs[param_name] = arg_type(context_id=context_id, dag_run_id=dag_run_id, default_ttl=default_ttl)
+            else:
+                func_kwargs[param_name] = args[pos_arg_index]
+                pos_arg_index += 1
+
+        result = task_func(**func_kwargs)
+
+        return result
+    return wrapper

plexflow/utils/tasks/k8s/task.py

@@ -0,0 +1,70 @@
+import json
+import os
+from airflow.providers.cncf.kubernetes.operators.pod import KubernetesPodOperator
+from kubernetes.client import models as k8s
+
+class PlexflowK8sOperator(KubernetesPodOperator):
+    def __init__(self, *args, **kwargs):
+        self.extra_args = args
+        self.extra_kwargs = kwargs.pop('extra_kwargs', {})
+        
+        # Define default volumes and volume mounts
+        self.default_volumes = [
+            k8s.V1Volume(
+                name='dags-volume',
+                persistent_volume_claim=k8s.V1PersistentVolumeClaimVolumeSource(claim_name='ssd-sync-dags-pvc'),
+            ),
+            k8s.V1Volume(
+                name='logs-volume',
+                persistent_volume_claim=k8s.V1PersistentVolumeClaimVolumeSource(claim_name='ssd-airflow-logs-pvc')
+            )
+        ]
+        self.default_volume_mounts = [
+            k8s.V1VolumeMount(
+                name='dags-volume',
+                mount_path='/dags',
+                sub_path='dags',
+            ),
+            k8s.V1VolumeMount(
+                name='logs-volume',
+                mount_path='/opt/airflow/logs'
+            )
+        ]
+        
+        # Add default volumes and volume mounts to the operator
+        if 'volumes' in kwargs:
+            kwargs['volumes'].extend(self.default_volumes)
+        else:
+            kwargs['volumes'] = self.default_volumes
+        
+        if 'volume_mounts' in kwargs:
+            kwargs['volume_mounts'].extend(self.default_volume_mounts)
+        else:
+            kwargs['volume_mounts'] = self.default_volume_mounts
+        
+        super().__init__(*args, **kwargs)
+
+    def execute(self, context):
+        labels = self._generate_labels(context)
+        env_vars = self._convert_labels_to_env_vars(labels)
+        self.env_vars.extend(env_vars)
+        super().execute(context)
+
+    def _generate_labels(self, context):
+        ti = context['ti']
+        labels = {
+            'TASK_MODE': 'k8s',
+            'AIRFLOW_DAG_ID': ti.dag_id,
+            'AIRFLOW_TASK_ID': ti.task_id,
+            'AIRFLOW_RUN_ID': context['run_id'],
+            'AIRFLOW_TRY_NUMBER': str(ti.try_number),
+            'AIRFLOW_MAX_TRIES': str(1 + ti.max_tries),
+            **{key.upper(): value for key, value in os.environ.items()},
+            **{f"ARG_{i:05d}": str(arg) for i, arg in enumerate(self.extra_args)},
+            **{f"KW_ARG_{i:05d}": json.dumps({"key": key, "value": value}) for i, (key, value) in enumerate(self.extra_kwargs.items())},
+        }
+        return labels
+
+    def _convert_labels_to_env_vars(self, labels):
+        # Converts labels to Kubernetes environment variables
+        return [k8s.V1EnvVar(name=k, value=v) for k, v in labels.items()]

plexflow/utils/thread_safe/safe_list.py

@@ -0,0 +1,54 @@
+import threading
+
+class ThreadSafeList:
+    def __init__(self):
+        self._list = []
+        self._lock = threading.Lock()
+
+    def append(self, item):
+        with self._lock:
+            self._list.append(item)
+
+    def remove(self, item):
+        with self._lock:
+            self._list.remove(item)
+
+    def __contains__(self, item):
+        with self._lock:
+            return item in self._list
+
+    def __getitem__(self, index):
+        with self._lock:
+            return self._list[index]
+
+    def __setitem__(self, index, value):
+        with self._lock:
+            self._list[index] = value
+
+    def __delitem__(self, index):
+        with self._lock:
+            del self._list[index]
+
+    def __iter__(self):
+        with self._lock:
+            return iter(self._list.copy())
+
+    def __len__(self):
+        with self._lock:
+            return len(self._list)
+
+    def insert(self, index, item):
+        with self._lock:
+            self._list.insert(index, item)
+    
+    def to_list(self):
+        with self._lock:
+            return list(self._list)
+    
+    def __str__(self) -> str:
+        with self._lock:
+            return str(self._list)
+    
+    def __repr__(self) -> str:
+        with self._lock:
+            return repr(self._list)

plexflow/utils/thread_safe/safe_set.py

@@ -0,0 +1,69 @@
+import threading
+
+class ThreadSafeSet:
+    def __init__(self):
+        self._set = set()
+        self._lock = threading.Lock()
+
+    def add(self, item):
+        with self._lock:
+            self._set.add(item)
+
+    def remove(self, item):
+        with self._lock:
+            self._set.remove(item)
+
+    def __contains__(self, item):
+        with self._lock:
+            return item in self._set
+
+    def __iter__(self):
+        with self._lock:
+            return iter(self._set.copy())
+
+    def __len__(self):
+        with self._lock:
+            return len(self._set)
+
+    def update(self, items):
+        with self._lock:
+            self._set.update(items)
+
+    def difference(self, other_set):
+        with self._lock:
+            if isinstance(other_set, ThreadSafeSet):
+                other_set = other_set._set
+            result = self._set.difference(other_set)
+        return ThreadSafeSet.from_set(result)
+
+    def intersection(self, other_set):
+        with self._lock:
+            if isinstance(other_set, ThreadSafeSet):
+                other_set = other_set._set
+            result = self._set.intersection(other_set)
+        return ThreadSafeSet.from_set(result)
+
+    def union(self, other_set):
+        with self._lock:
+            if isinstance(other_set, ThreadSafeSet):
+                other_set = other_set._set
+            result = self._set.union(other_set)
+        return ThreadSafeSet.from_set(result)
+    
+    def to_set(self):
+        with self._lock:
+            return set(self._set)
+    
+    def __str__(self) -> str:
+        with self._lock:
+            return str(self._set)
+    
+    def __repr__(self) -> str:
+        with self._lock:
+            return repr(self._set)
+    
+    @classmethod
+    def from_set(cls, input_set):
+        instance = cls()
+        instance._set = input_set
+        return instance

plexflow/utils/torrent/analyze.py

@@ -0,0 +1,118 @@
+from plexflow.core.torrents.results.torrent import Torrent
+from plexflow.utils.imdb.imdb_codes import IMDbCode, extract_imdb_code
+import requests
+import logging
+from plexflow.utils.subtitle.search import SubtitleSearcher
+from bs4 import BeautifulSoup
+
+class TorrentReport:
+    def __init__(self, **kwargs) -> None:
+        self._torrent: Torrent = kwargs.get("torrent")
+        self._extracted_imdb_code = kwargs.get("extracted_imdb_code")
+        self.hardcoded = kwargs.get("hardcoded")
+        self.korsub = kwargs.get("korsub")
+        self._subtitles = kwargs.get("subtitles")
+    
+    @property
+    def torrent(self) -> Torrent:
+        return self._torrent
+    
+    @property
+    def extracted_imdb_code(self) -> str:
+        return self._extracted_imdb_code
+    
+    @property
+    def imdb_code_matched(self) -> bool:
+        return IMDbCode(self.torrent.imdb_code) == IMDbCode(self.extracted_imdb_code)
+
+    @property
+    def acceptable_quality(self) -> bool:
+        return self.torrent.parsed_release_name.get("quality", "").upper() not in [
+            "CAM", "TS", "TC", "SCR", "DVDSCR",
+            "SCREENER", "TELESYNC", "TELECINE", "DVDSCREENER",
+            "BDSCR", "WEBSCREENER", "HDCAM",
+        ]
+    @property
+    def has_hardcoded_subtitles(self) -> bool:
+        return self.torrent.parsed_release_name.get("hardcoded", False) or self.hardcoded
+    
+    @property
+    def has_korsub_subtitles(self) -> bool:
+        return self.korsub
+    
+    @property
+    def subtitles(self) -> list:
+        return self._subtitles
+
+class TorrentInspector:
+    def __init__(self, torrent: Torrent) -> None:
+        self.torrent = torrent
+    
+    def inspect(self) -> TorrentReport:
+        report = {
+            "torrent": self.torrent,
+        }
+        
+        logging.info(f"Inspecting torrent: {self.torrent}")
+        logging.info(f"Inspecting release name: {self.torrent.release_name}")
+        logging.info(f"Inspecting IMDb code: {self.torrent.imdb_code}")
+        logging.info(f"Inspecting URL: {self.torrent.url}")
+        
+        try:
+            url = self.torrent.url
+            if isinstance(url, str) and len(url) > 0:
+                response = requests.get(
+                    url=url,
+                    headers={
+                        "accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.7",
+                    }
+                )
+                response.raise_for_status()
+
+                logging.info(f"URL status code: {response.status_code}")
+                
+                extracted_imdb_id = next(extract_imdb_code(response.text), None)
+                logging.info(f"Extracted IMDb code: {extracted_imdb_id}")
+                
+                soup = BeautifulSoup(response.text, 'html.parser')
+                
+                report["extracted_imdb_code"] = extracted_imdb_id
+
+                # check if torrent has hardcoded subtitles using various alternatives for the word
+                # hardcoded
+                hardcoded = any([
+                    "hardcoded" in self.torrent.release_name.lower(),
+                    "hardsub" in self.torrent.release_name.lower(),
+                    "hardcoded" in response.text.lower(),
+                    "hardsub" in response.text.lower(),
+                ])
+                
+                logging.info(f"Hardcoded subtitles: {hardcoded}")
+                
+                # check if torrent has korsub subtitles
+                korsub = any([
+                    "korsub" in self.torrent.release_name.lower(),
+                    "korsub" in response.text.lower(),
+                ])
+                
+                logging.info(f"Korsub subtitles: {korsub}")
+                
+                report["hardcoded"] = hardcoded
+                report["korsub"] = korsub
+                
+                searcher = SubtitleSearcher(hint_words=[
+                    "english",
+                    "eng",
+                    "dutch",
+                    "dut",
+                ])
+                
+                subtitles = searcher.search_subtitles(soup.get_text())
+                logging.info(f"Subtitles found: {subtitles}")
+                report["subtitles"] = subtitles.split(",")
+            else:
+                logging.info("No URL provided for torrent")
+        except Exception as e:
+            logging.error(f"Error while inspecting torrent: {e}")
+
+        return TorrentReport(**report)

plexflow/utils/torrent/extract/common.py

@@ -0,0 +1,37 @@
+from plexflow.utils.subtitle.search import SubtitleSearcher
+from bs4 import BeautifulSoup
+from plexflow.utils.imdb.imdb_codes import extract_imdb_code
+from plexflow.utils.torrent.hash import extract_magnet
+from plexflow.utils.torrent.files import TorrentSubtitle
+from plexflow.utils.strings.language import get_language_code
+
+def torrent_detail_extract(func):
+    def wrapper(html):
+        searcher = SubtitleSearcher(
+            hint_words={
+                "dutch",
+                "dut",
+                "eng",
+                "english",
+            }
+        )
+
+        soup = BeautifulSoup(html, 'html.parser')
+        text = soup.get_text()
+        
+        subtitles = searcher.search_subtitles(text)
+        
+        imdb_code = next(extract_imdb_code(html), None)
+
+        magnet = next(extract_magnet(text), None)
+
+        details = {
+            "subtitles": [TorrentSubtitle(name=sub, language=get_language_code(sub)) for sub in filter(lambda s: s.strip(), subtitles.split(","))],
+            "imdb_id": imdb_code,
+            "magnet": magnet,
+        }
+
+        # Call the decorated function with the HTML and extracted parts as a dictionary
+        return func(html, **details)
+
+    return wrapper