Refactor pre-queue script

[skip ci]

sabnzbd/newsunpack.py

@@ -36,7 +36,6 @@ import sabnzbd.utils.rarfile as rarfile
 from sabnzbd.misc import (
     format_time_string,
     find_on_path,
-    int_conv,
     get_all_passwords,
     calc_age,
     cmp,
@@ -45,7 +44,6 @@ from sabnzbd.misc import (
     format_time_left,
 )
 from sabnzbd.filesystem import (
-    make_script_path,
     real_path,
     globber,
     globber_full,
@@ -2134,89 +2132,6 @@ def add_time_left(perc: float, start_time: Optional[float] = None, time_used: Op
         return " - %s %s" % (format_time_left(int((100 - perc) / (perc / time_used)), short_format=True), T("left"))
     return ""
 
-
-def pre_queue(nzo: NzbObject, pp, cat):
-    """Run pre-queue script (if any) and process results.
-    pp and cat are supplied separate since they can change.
-    """
-
-    def fix(p):
-        # If added via API, some items can still be "None" (as a string)
-        if not p or str(p).lower() == "none":
-            return ""
-        return str(p)
-
-    values = [1, nzo.final_name_with_password, pp, cat, nzo.script, nzo.priority, None]
-    script_path = make_script_path(cfg.pre_script())
-    if script_path:
-        # Basic command-line parameters
-        command = [
-            script_path,
-            nzo.final_name_with_password,
-            pp,
-            cat,
-            nzo.script,
-            nzo.priority,
-            str(nzo.bytes),
-            " ".join(nzo.groups),
-        ]
-        command.extend(list(sabnzbd.sorting.analyse_show(nzo.final_name).values()))
-        command = [fix(arg) for arg in command]
-
-        # Fields not in the NZO directly
-        extra_env_fields = {
-            "groups": " ".join(nzo.groups),
-            "show_name": command[8],
-            "show_season": command[9],
-            "show_episode": command[10],
-            "show_episode_name": command[11],
-            "proper": command[12],
-            "resolution": command[13],
-            "decade": command[14],
-            "year": command[15],
-            "month": command[16],
-            "day": command[17],
-            "type": command[18],
-        }
-
-        try:
-            p = build_and_run_command(command, env=create_env(nzo, extra_env_fields))
-        except:
-            logging.debug("Failed script %s, Traceback: ", script_path, exc_info=True)
-            return values
-
-        output = p.stdout.read()
-        ret = p.wait()
-        logging.info("Pre-queue script returned %s and output=\n%s", ret, output)
-        if ret == 0:
-            split_output = output.splitlines()
-            try:
-                # Extract category line from pre-queue output
-                pre_queue_category = split_output[3].strip(" '\"")
-            except IndexError:
-                pre_queue_category = None
-
-            for index, line in enumerate(split_output):
-                line = line.strip(" '\"")
-                if index < len(values):
-                    if line:
-                        values[index] = line
-                    elif pre_queue_category and index in (2, 4, 5):
-                        # Preserve empty pp, script, and priority lines to prevent
-                        # pre-existing values from overriding category-based settings
-                        values[index] = ""
-
-        accept = int_conv(values[0])
-        if accept < 1:
-            logging.info("Pre-Q refuses %s", nzo.final_name)
-        elif accept == 2:
-            logging.info("Pre-Q accepts&fails %s", nzo.final_name)
-        else:
-            logging.info("Pre-Q accepts %s", nzo.final_name)
-
-    return values
-
-
 def is_sevenfile(path: str) -> bool:
     """Return True if path has 7Zip-signature and 7Zip is detected"""
     with open(path, "rb") as sevenzip:

sabnzbd/nzbstuff.py

@@ -826,59 +826,15 @@ class NzbObject(TryList):
         # Determine category and find pp/script values
         self.cat, pp_tmp, self.script, priority = cat_to_opts(cat, pp, script, self.priority)
         self.set_priority(priority)
-        self.repair, self.unpack, self.delete = pp_to_opts(pp_tmp)
+        self.set_pp(pp_tmp)
 
         # Show first meta-password (if any), when there's no explicit password
         if not self.password and self.meta.get("password"):
             self.password = self.meta.get("password", [None])[0]
 
-        # Run user pre-queue script if set and valid
+        # Run user pre-queue script
-        if not reuse and make_script_path(cfg.pre_script()):
+        if not reuse:
-            # Call the script
+            self.run_pre_queue(input_priority, input_pp, input_script)
-            accept, name, pq_pp, pq_cat, pq_script, pq_priority, pq_group = sabnzbd.newsunpack.pre_queue(self, pp, cat)
-
-            if pq_cat:
-                # An explicit pp/script/priority set upon adding the job takes precedence
-                # over an implicit setting based on the category set by pre-queue
-                if input_priority and not pq_priority:
-                    pq_priority = input_priority
-                if input_pp and not pq_pp:
-                    pq_pp = input_pp
-                if input_script and not pq_script:
-                    pq_script = input_script
-
-            # Accept or reject
-            accept = int_conv(accept)
-            if accept < 1:
-                self.purge_data()
-                raise NzbRejected
-            if accept == 2:
-                raise NzbRejectToHistory(self, T("Pre-queue script marked job as failed"))
-
-            # Process all options, only over-write if set by script
-            # Beware that cannot do "if priority/pp", because those can
-            # also have a valid value of 0, which shouldn't be ignored
-            if name:
-                self.set_final_name_and_scan_password(name)
-            try:
-                pp = int(pq_pp)
-            except:
-                pp = None
-            if pq_cat:
-                cat = pq_cat
-            try:
-                priority = int(pq_priority)
-            except:
-                priority = DEFAULT_PRIORITY
-            if pq_script and is_valid_script(pq_script):
-                script = pq_script
-            if pq_group:
-                self.groups = [str(pq_group)]
-
-            # Re-evaluate results from pre-queue script
-            self.cat, pp, self.script, priority = cat_to_opts(cat, pp, script, priority)
-            self.set_priority(priority)
-            self.repair, self.unpack, self.delete = pp_to_opts(pp)
 
         # Pause if requested by the NZB-adding or the pre-queue script
         if self.priority == PAUSED_PRIORITY:
@@ -1981,6 +1937,96 @@ class NzbObject(TryList):
             logging.info("Pausing duplicate alternative %s", self.final_name)
             self.pause()
 
+    def run_pre_queue(
+        self,
+        input_priority: Optional[Union[int, str]],
+        input_pp: Optional[int],
+        input_script: Optional[str],
+    ):
+        """Run pre-queue script (if any) and process results."""
+        if script_path := make_script_path(cfg.pre_script()):
+
+            def fix_parameter(parameter: Any) -> str:
+                # If added via API, some items can still be "None" (as a string)
+                if not parameter or str(parameter).lower() == "none":
+                    return ""
+                return str(parameter)
+
+            # Basic parameters
+            command = [script_path, self.final_name_with_password, self.cat, self.priority, self.pp, self.script]
+            command = [fix_parameter(arg) for arg in command]
+
+            # Fields not in the NZO directly
+            extra_env_fields = sabnzbd.newsunpack.analyse_show(self.final_name_with_password)
+            extra_env_fields["groups"] = " ".join(self.groups)
+
+            try:
+                p = sabnzbd.newsunpack.build_and_run_command(
+                    command, env=sabnzbd.newsunpack.create_env(self, extra_env_fields)
+                )
+            except:
+                logging.debug("Failed script %s, Traceback: ", script_path, exc_info=True)
+                return
+
+            output = p.stdout.read()
+            ret = p.wait()
+            logging.info("Pre-queue script returned %s and output=\n%s", ret, output)
+            if ret == 0:
+                # Base values
+                pq_cat = pq_pp = pq_script = pq_priority = None
+                for index, line in enumerate(output.splitlines(), start=1):
+                    # Make sure to keep this in line with the documentation!
+                    # 1: Accept
+                    # 2: Name
+                    # 3: Category
+                    # 4: Priority
+                    # 5: Post-processing
+                    # 6: Script
+                    # 7: Duplicate
+                    # 8: Duplicate key
+                    if line := line.strip(" '\""):
+                        if index == 1:
+                            # Accept or reject
+                            accept = int_conv(line)
+                            if accept < 1:
+                                logging.info("Pre-queue script refuses %s", self.final_name)
+                                self.purge_data()
+                                raise NzbRejected
+                            if accept == 2:
+                                logging.info("Pre-queue marking as failed %s", self.final_name)
+                                raise NzbRejectToHistory(self, T("Pre-queue script marked job as failed"))
+                            logging.info("Pre-queue accepts %s", self.final_name)
+                        elif index == 2:
+                            self.set_final_name_and_scan_password(line)
+                        elif index == 3:
+                            pq_cat = line
+                        elif index == 4:
+                            pq_priority = int_conv(line, default=DEFAULT_PRIORITY)
+                        elif index == 5:
+                            pq_pp = int_conv(line, default=None)
+                        elif index == 6:
+                            if is_valid_script(line):
+                                pq_script = line
+                        elif index == 7:
+                            self.duplicate = line
+                        elif index == 8:
+                            self.duplicate_series_key = line
+
+                if pq_cat:
+                    # An explicit pp/script/priority set upon adding the job takes precedence
+                    # over an implicit setting based on the category set by pre-queue
+                    if input_priority and pq_priority is None:
+                        pq_priority = input_priority
+                    if input_pp and pq_pp is None:
+                        pq_pp = input_pp
+                    if input_script and not pq_script:
+                        pq_script = input_script
+
+                # Re-evaluate results from pre-queue script
+                self.cat, pp, self.script, priority = cat_to_opts(pq_cat, pq_pp, pq_script, pq_priority)
+                self.set_priority(priority)
+                self.set_pp(pp)
+
     def __getstate__(self):
         """Save to pickle file, selecting attributes"""
         dict_ = {}

tests/test_functional_adding_nzbs.py

@@ -174,9 +174,17 @@ class TestAddingNZBs:
         try:
             script_path = os.path.join(VAR.SCRIPT_DIR, script_name)
             with open(script_path, "w") as f:
-                # line 1 = accept; 4 = category; 6 = priority
+                # Lines:
+                # 1: Accept
+                # 2: Name
+                # 3: Category
+                # 4: Priority
+                # 5: Post-processing
+                # 6: Script
+                # 7: Duplicate
+                # 8: Duplicate key
                 f.write(
-                    "#!%s\n\nprint('1\\n\\n\\n%s\\n\\n%s\\n')"
+                    "#!%s\n\nprint('1\\n\\n%s\\n%s\\n')"
                     % (
                         sys.executable,
                         (category if category else ""),
@@ -406,8 +414,8 @@ class TestAddingNZBs:
     @pytest.mark.parametrize("prio_def_cat", sample(VALID_DEFAULT_PRIORITIES, 2))
     @pytest.mark.parametrize("prio_add", sample(PRIO_OPTS_ADD, 3))
     @pytest.mark.parametrize("prio_add_cat", sample(PRIO_OPTS_ADD_CAT, 2))
-    @pytest.mark.parametrize("prio_preq", sample(PRIO_OPTS_PREQ, 2))
+    @pytest.mark.parametrize("prio_preq", PRIO_OPTS_PREQ)
-    @pytest.mark.parametrize("prio_preq_cat", sample(PRIO_OPTS_PREQ_CAT, 2))
+    @pytest.mark.parametrize("prio_preq_cat", PRIO_OPTS_PREQ_CAT)
     def test_adding_nzbs_priority_sample(
         self, prio_def_cat, prio_add, prio_add_cat, prio_preq, prio_preq_cat, prio_meta_cat
     ):