MAINT: add threads/jobs primitives (#82)

qiime2 · Feb 13, 2024 · a0e44f7 · a0e44f7
1 parent f4a4883
commit a0e44f7
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 5 deletions.
diff --git a/q2_alignment/_mafft.py b/q2_alignment/_mafft.py
@@ -79,7 +79,7 @@ def _mafft(sequences_fp, alignment_fp, n_threads, parttree, addfragments):
     # mafft's signal for utilizing all cores is -1. We want to our users
     # to enter auto for using all cores. This is to prevent any confusion and
     # to keep the UX consisent.
-    if n_threads == 'auto':
+    if n_threads == 0:
         n_threads = -1
 
     # `--inputorder` must be turned on because we need the input and output in

diff --git a/q2_alignment/plugin_setup.py b/q2_alignment/plugin_setup.py
@@ -7,7 +7,7 @@
 # ----------------------------------------------------------------------------
 
 from qiime2.plugin import (
-    Plugin, Float, Int, Bool, Range, Citations, Str, Choices)
+    Plugin, Float, Bool, Range, Citations, Threads)
 from q2_types.feature_data import FeatureData, Sequence, AlignedSequence
 
 import q2_alignment
@@ -26,7 +26,7 @@
 plugin.methods.register_function(
     function=q2_alignment.mafft,
     inputs={'sequences': FeatureData[Sequence]},
-    parameters={'n_threads': Int % Range(1, None) | Str % Choices(['auto']),
+    parameters={'n_threads': Threads,
                 'parttree': Bool},
     outputs=[('alignment', FeatureData[AlignedSequence])],
     input_descriptions={'sequences': 'The sequences to be aligned.'},
@@ -45,7 +45,7 @@
     function=q2_alignment.mafft_add,
     inputs={'alignment': FeatureData[AlignedSequence],
             'sequences': FeatureData[Sequence]},
-    parameters={'n_threads': Int % Range(1, None) | Str % Choices(['auto']),
+    parameters={'n_threads': Threads,
                 'parttree': Bool,
                 'addfragments': Bool},
     outputs=[('expanded_alignment', FeatureData[AlignedSequence])],

diff --git a/q2_alignment/tests/test_mafft.py b/q2_alignment/tests/test_mafft.py
@@ -45,7 +45,7 @@ def test_multithreaded_mafft(self):
         input_sequences, exp = self._prepare_sequence_data()
 
         with redirected_stdio(stderr=os.devnull):
-            result = mafft(input_sequences, n_threads='auto')
+            result = mafft(input_sequences, n_threads=0)
         obs = skbio.io.read(str(result), into=skbio.TabularMSA,
                             constructor=skbio.DNA)
         self.assertEqual(obs, exp)