Add support for tagging inputs in the run management command

Signed-off-by: tdruez <[email protected]>

Author: tdruez

scanpipe/management/commands/__init__.py

@@ -284,20 +284,23 @@ def validate_pipelines(pipelines_data):
     return pipelines_data
 
 
-def extract_tag_from_input_files(input_files):
+def extract_tag_from_input_file(file_location):
     """
-    Add support for the ":tag" suffix in file location.
+    Parse a file location with optional tag suffix.
 
     For example: "/path/to/file.zip:tag"
     """
-    input_files_data = {}
-    for file in input_files:
-        if ":" in file:
-            key, value = file.split(":", maxsplit=1)
-            input_files_data.update({key: value})
-        else:
-            input_files_data.update({file: ""})
-    return input_files_data
+    if ":" in file_location:
+        cleaned_location, tag = file_location.split(":", maxsplit=1)
+        return cleaned_location, tag
+    return file_location, ""
+
+
+def extract_tag_from_input_files(input_files):
+    """Parse multiple file locations with optional tag suffixes."""
+    return dict(
+        extract_tag_from_input_file(file_location) for file_location in input_files
+    )
 
 
 def validate_input_files(input_files):

scanpipe/management/commands/run.py

@@ -28,6 +28,7 @@
 from django.core.management.base import CommandError
 from django.utils.crypto import get_random_string
 
+from scanpipe.management.commands import extract_tag_from_input_file
 from scanpipe.pipes.fetch import SCHEME_TO_FETCHER_MAPPING
 
 
@@ -94,8 +95,10 @@ def get_input_options(input_location):
         for location in input_location.split(","):
             if location.startswith(tuple(SCHEME_TO_FETCHER_MAPPING.keys())):
                 input_options["input_urls"].append(location)
+
             else:
-                input_path = Path(location)
+                cleaned_location, _ = extract_tag_from_input_file(location)
+                input_path = Path(cleaned_location)
                 if not input_path.exists():
                     raise CommandError(f"{location} not found.")
                 if input_path.is_file():

scanpipe/tests/test_commands.py

@@ -1461,6 +1461,19 @@ def test_scanpipe_management_command_verify_project(self):
                 stdout=out,
             )
 
+    def test_scanpipe_management_command_extract_tag_from_input_file(self):
+        extract_tag = commands.extract_tag_from_input_file
+        expected = ("file.ext", "")
+        self.assertEqual(expected, extract_tag("file.ext"))
+        expected = ("file.ext", "")
+        self.assertEqual(expected, extract_tag("file.ext:"))
+        expected = ("file.ext", "tag")
+        self.assertEqual(expected, extract_tag("file.ext:tag"))
+        expected = ("file.ext", "tag1:tag2")
+        self.assertEqual(expected, extract_tag("file.ext:tag1:tag2"))
+        expected = ("file.ext", "tag1,tag2")
+        self.assertEqual(expected, extract_tag("file.ext:tag1,tag2"))
+
 
 class ScanPipeManagementCommandMixinTest(TestCase):
     class CreateProjectCommand(