From a895429948be9477eab681763bc36dfab3fbbbed Mon Sep 17 00:00:00 2001
From: VinzentRisch <100149044+VinzentRisch@users.noreply.github.com>
Date: Fri, 26 Jul 2024 13:39:14 +0200
Subject: [PATCH] ENH: Add three new parameters to run_amrfinderplus_n (#97)

---
 q2_amr/amrfinderplus/sample_data.py           |  5 ++
 .../amrfinderplus/tests/test_sample_data.py   |  3 +
 q2_amr/amrfinderplus/tests/test_utils.py      | 56 +++++++++++++++++++
 q2_amr/amrfinderplus/utils.py                 | 16 ++++++
 q2_amr/plugin_setup.py                        | 43 +++++++++++++-
 5 files changed, 122 insertions(+), 1 deletion(-)

diff --git a/q2_amr/amrfinderplus/sample_data.py b/q2_amr/amrfinderplus/sample_data.py
index 9a7c4fa..06de759 100644
--- a/q2_amr/amrfinderplus/sample_data.py
+++ b/q2_amr/amrfinderplus/sample_data.py
@@ -25,6 +25,8 @@ def annotate_sample_data_amrfinderplus(
     curated_ident: bool = False,
     coverage_min: float = 0.5,
     translation_table: str = "11",
+    report_common: bool = False,
+    gpipe_org: bool = False,
     threads: int = None,
 ) -> (
     AMRFinderPlusAnnotationsDirFmt,
@@ -72,6 +74,9 @@ def annotate_sample_data_amrfinderplus(
                 curated_ident=curated_ident,
                 coverage_min=coverage_min,
                 translation_table=translation_table,
+                annotation_format=None,
+                report_common=report_common,
+                gpipe_org=gpipe_org,
                 threads=threads,
             )
 
diff --git a/q2_amr/amrfinderplus/tests/test_sample_data.py b/q2_amr/amrfinderplus/tests/test_sample_data.py
index 8c2ff05..9d4557b 100644
--- a/q2_amr/amrfinderplus/tests/test_sample_data.py
+++ b/q2_amr/amrfinderplus/tests/test_sample_data.py
@@ -25,6 +25,9 @@ def mock_run_amrfinderplus_n(
         curated_ident,
         coverage_min,
         translation_table,
+        annotation_format,
+        report_common,
+        gpipe_org,
         threads,
     ):
         with open(os.path.join(working_dir, "amr_annotations.tsv"), "w"):
diff --git a/q2_amr/amrfinderplus/tests/test_utils.py b/q2_amr/amrfinderplus/tests/test_utils.py
index 4b7f436..0bd6b73 100644
--- a/q2_amr/amrfinderplus/tests/test_utils.py
+++ b/q2_amr/amrfinderplus/tests/test_utils.py
@@ -23,6 +23,9 @@ def test_run_amrfinderplus_n(self, mock_run_command):
             curated_ident=False,
             coverage_min=1,
             translation_table="11",
+            annotation_format="prodigal",
+            report_common=True,
+            gpipe_org=True,
             threads=4,
         )
         mock_run_command.assert_called_once_with(
@@ -57,6 +60,10 @@ def test_run_amrfinderplus_n(self, mock_run_command):
                 "1",
                 "--translation_table",
                 "11",
+                "--annotation_format",
+                "prodigal",
+                "--report_common",
+                "--gpipe_org",
             ],
             "path_dir",
             verbose=True,
@@ -77,6 +84,9 @@ def test_run_amrfinderplus_n_minimal(self, mock_run_command):
             curated_ident=True,
             coverage_min=None,
             translation_table=None,
+            annotation_format=None,
+            report_common=False,
+            gpipe_org=False,
             threads=None,
         )
         mock_run_command.assert_called_once_with(
@@ -93,3 +103,49 @@ def test_run_amrfinderplus_n_minimal(self, mock_run_command):
             "path_dir",
             verbose=True,
         )
+
+    @patch("q2_amr.amrfinderplus.utils.run_command")
+    def test_run_amrfinderplus_n_value_error_report_common(self, mock_run_command):
+        with self.assertRaisesRegex(
+            ValueError, "--p-report-common requires " "--p-plus and --p-organism"
+        ):
+            run_amrfinderplus_n(
+                working_dir="path_dir",
+                amrfinderplus_db="amrfinderplus_db",
+                dna_sequences=None,
+                protein_sequences=None,
+                gff=None,
+                organism=None,
+                plus=False,
+                report_all_equal=False,
+                ident_min=None,
+                curated_ident=True,
+                coverage_min=None,
+                translation_table=None,
+                annotation_format=None,
+                report_common=True,
+                gpipe_org=False,
+                threads=None,
+            )
+
+    @patch("q2_amr.amrfinderplus.utils.run_command")
+    def test_run_amrfinderplus_n_value_error_gpipe_org(self, mock_run_command):
+        with self.assertRaisesRegex(ValueError, "--p-gpipe_org requires --p-organism"):
+            run_amrfinderplus_n(
+                working_dir="path_dir",
+                amrfinderplus_db="amrfinderplus_db",
+                dna_sequences=None,
+                protein_sequences=None,
+                gff=None,
+                organism=None,
+                plus=False,
+                report_all_equal=False,
+                ident_min=None,
+                curated_ident=True,
+                coverage_min=None,
+                translation_table=None,
+                annotation_format=None,
+                report_common=False,
+                gpipe_org=True,
+                threads=None,
+            )
diff --git a/q2_amr/amrfinderplus/utils.py b/q2_amr/amrfinderplus/utils.py
index 793ee37..51dfcb3 100644
--- a/q2_amr/amrfinderplus/utils.py
+++ b/q2_amr/amrfinderplus/utils.py
@@ -16,8 +16,17 @@ def run_amrfinderplus_n(
     curated_ident,
     coverage_min,
     translation_table,
+    annotation_format,
+    report_common,
+    gpipe_org,
     threads,
 ):
+    # Check for unallowed parameter combinations
+    if report_common and (not plus or not organism):
+        raise ValueError("--p-report-common requires --p-plus and --p-organism")
+    if gpipe_org and not organism:
+        raise ValueError("--p-gpipe_org requires --p-organism")
+
     cmd = [
         "amrfinder",
         "--database",
@@ -73,6 +82,13 @@ def run_amrfinderplus_n(
         cmd.extend(["--coverage_min", str(coverage_min)])
     if translation_table:
         cmd.extend(["--translation_table", str(translation_table)])
+    if annotation_format:
+        cmd.extend(["--annotation_format", str(annotation_format)])
+    if report_common:
+        cmd.append("--report_common")
+    if gpipe_org:
+        cmd.append("--gpipe_org")
+
     try:
         run_command(cmd, working_dir, verbose=True)
     except subprocess.CalledProcessError as e:
diff --git a/q2_amr/plugin_setup.py b/q2_amr/plugin_setup.py
index 703cecb..b57aeb6 100644
--- a/q2_amr/plugin_setup.py
+++ b/q2_amr/plugin_setup.py
@@ -1127,6 +1127,35 @@
     "Vibrio_vulnificus",
 ]
 
+organisms_gpipe = [
+    "Acinetobacter",
+    "Burkholderia_cepacia_complex",
+    "Burkholderia_pseudomallei",
+    "Campylobacter",
+    "Citrobacter_freundii",
+    "Clostridioides_difficile",
+    "Enterobacter_asburiae",
+    "Enterobacter_cloacae",
+    "Enterococcus_faecalis",
+    "Enterococcus_faecium",
+    "Escherichia_coli_Shigella",
+    "Klebsiella_oxytoca",
+    "Klebsiella",
+    "Neisseria_gonorrhoeae",
+    "Neisseria_meningitidis",
+    "Pseudomonas_aeruginosa",
+    "Salmonella",
+    "Serratia",
+    "Staphylococcus_aureus",
+    "Staphylococcus_pseudintermedius",
+    "Streptococcus_agalactiae",
+    "Streptococcus_pneumoniae",
+    "Streptococcus_pyogenes",
+    "Vibrio_cholerae",
+    "Vibrio_parahaemolyticus",
+    "Vibrio_vulnificus",
+]
+
 translation_tables = [
     "1",
     "2",
@@ -1156,6 +1185,13 @@
     "33",
 ]
 
+P_gpipe_org, P_organism, _ = TypeMap(
+    {
+        (Bool % Choices(True), Str % Choices(organisms_gpipe)): Int,
+        (Bool % Choices(False), Str % Choices(organisms)): Int,
+    }
+)
+
 plugin.methods.register_function(
     function=annotate_sample_data_amrfinderplus,
     inputs={
@@ -1163,13 +1199,15 @@
         "amrfinderplus_db": AMRFinderPlusDatabase,
     },
     parameters={
-        "organism": Str % Choices(organisms),
+        "organism": P_organism,
         "plus": Bool,
         "report_all_equal": Bool,
         "ident_min": Float % Range(0, 1, inclusive_start=True, inclusive_end=True),
         "curated_ident": Bool,
         "coverage_min": Float % Range(0, 1, inclusive_start=True, inclusive_end=True),
         "translation_table": Str % Choices(translation_tables),
+        "report_common": Bool,
+        "gpipe_org": P_gpipe_org,
         "threads": Int % Range(0, None, inclusive_start=False),
     },
     outputs=[
@@ -1204,6 +1242,9 @@
         "coverage_min": "Minimum proportion of reference gene covered for a "
         "BLAST-based hit (Methods BLAST or PARTIAL).",
         "translation_table": "Translation table used for BLASTX.",
+        "report_common": "Report proteins common to a taxonomy group.",
+        "gpipe_org": "Use Pathogen Detection taxgroup names as arguments to the "
+        "organism option.",
         "threads": "The number of threads to use for processing. AMRFinderPlus "
         "defaults to 4 on hosts with >= 4 cores. Setting this number higher"
         " than the number of cores on the running host may cause blastp to "