adapt action with new type

caporaso-lab · Apr 6, 2024 · da9e5bf · da9e5bf
1 parent 4bfb642
commit da9e5bf
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 41 deletions.
diff --git a/q2_dwq2/_methods.py b/q2_dwq2/_methods.py
@@ -7,19 +7,15 @@
 # ----------------------------------------------------------------------------
 
 from skbio.alignment import global_pairwise_align_nucleotide, TabularMSA
+from skbio import DNA
 
-from q2_types.feature_data import DNAIterator
 
-
-def nw_align(seq1: DNAIterator,
-             seq2: DNAIterator,
+def nw_align(seq1: DNA,
+             seq2: DNA,
              gap_open_penalty: float = 5,
              gap_extend_penalty: float = 2,
              match_score: float = 1,
              mismatch_score: float = -2) -> TabularMSA:
-    seq1 = next(iter(seq1))
-    seq2 = next(iter(seq2))
-
     msa, _, _ = global_pairwise_align_nucleotide(
         seq1=seq1, seq2=seq2, gap_open_penalty=gap_open_penalty,
         gap_extend_penalty=gap_extend_penalty, match_score=match_score,

diff --git a/q2_dwq2/plugin_setup.py b/q2_dwq2/plugin_setup.py
@@ -9,7 +9,7 @@
 import importlib
 
 from qiime2.plugin import Citations, Plugin, Float, Range
-from q2_types.feature_data import FeatureData, Sequence, AlignedSequence
+from q2_types.feature_data import FeatureData, AlignedSequence
 from q2_dwq2 import __version__
 from q2_dwq2._methods import nw_align
 from q2_dwq2._visualizers import summarize_alignment
@@ -44,8 +44,8 @@
 # Register actions
 plugin.methods.register_function(
     function=nw_align,
-    inputs={'seq1': FeatureData[Sequence],
-            'seq2': FeatureData[Sequence]},
+    inputs={'seq1': SingleDNASequence,
+            'seq2': SingleDNASequence},
     parameters={
         'gap_open_penalty': Float % Range(0, None, inclusive_start=False),
         'gap_extend_penalty': Float % Range(0, None, inclusive_start=False),

diff --git a/q2_dwq2/tests/test_methods.py b/q2_dwq2/tests/test_methods.py
@@ -11,9 +11,9 @@
 
 from qiime2.plugin.testing import TestPluginBase
 from qiime2.plugin.util import transform
-from q2_types.feature_data import DNAFASTAFormat, DNAIterator
 
 from q2_dwq2._methods import nw_align
+from q2_dwq2._types_and_formats import SingleRecordDNAFASTAFormat
 
 
 class NWAlignTests(TestPluginBase):
@@ -22,9 +22,7 @@ class NWAlignTests(TestPluginBase):
     def test_simple1(self):
         # test alignment of a pair of sequences
         sequence1 = DNA('AAAAAAAAGGTGGCCTTTTTTTT')
-        sequence1 = DNAIterator([sequence1])
         sequence2 = DNA('AAAAAAAAGGGGCCTTTTTTTT')
-        sequence2 = DNAIterator([sequence2])
         observed = nw_align(sequence1, sequence2)
 
         aligned_sequence1 = DNA('AAAAAAAAGGTGGCCTTTTTTTT')
@@ -38,12 +36,12 @@ def test_simple2(self):
         # loaded from file this time, for demonstration purposes
         sequence1 = transform(
             self.get_data_path('seq-1.fasta'),
-            from_type=DNAFASTAFormat,
-            to_type=DNAIterator)
+            from_type=SingleRecordDNAFASTAFormat,
+            to_type=DNA)
         sequence2 = transform(
             self.get_data_path('seq-2.fasta'),
-            from_type=DNAFASTAFormat,
-            to_type=DNAIterator)
+            from_type=SingleRecordDNAFASTAFormat,
+            to_type=DNA)
         observed = nw_align(sequence1, sequence2)
 
         aligned_sequence1 = DNA('ACCGGTGGAACCGG-TAACACCCAC')
@@ -53,10 +51,8 @@ def test_simple2(self):
         self.assertNotEqual(observed, expected)
 
     def test_alt_match_score(self):
-        s1 = DNA('AAAATTT')
-        sequence1 = DNAIterator([s1])
-        s2 = DNA('AAAAGGTTT')
-        sequence2 = DNAIterator([s2])
+        sequence1 = DNA('AAAATTT')
+        sequence2 = DNA('AAAAGGTTT')
         # call with default value for match score
         observed = nw_align(sequence1, sequence2)
 
@@ -66,8 +62,6 @@ def test_alt_match_score(self):
 
         self.assertEqual(observed, expected)
 
-        sequence1 = DNAIterator([s1])
-        sequence2 = DNAIterator([s2])
         # call with non-default value for match_score
         observed = nw_align(sequence1, sequence2, match_score=10)
 
@@ -86,10 +80,8 @@ def test_alt_match_score(self):
         self.assertEqual(observed, expected)
 
     def test_alt_gap_open_penalty(self):
-        s1 = DNA('AAAATTT')
-        sequence1 = DNAIterator([s1])
-        s2 = DNA('AAAAGGTTT')
-        sequence2 = DNAIterator([s2])
+        sequence1 = DNA('AAAATTT')
+        sequence2 = DNA('AAAAGGTTT')
         observed = nw_align(sequence1, sequence2, gap_open_penalty=0.01)
 
         aligned_sequence1 = DNA('AAAA-T-TT-')
@@ -98,8 +90,6 @@ def test_alt_gap_open_penalty(self):
 
         self.assertEqual(observed, expected)
 
-        sequence1 = DNAIterator([s1])
-        sequence2 = DNAIterator([s2])
         observed = nw_align(sequence1, sequence2)
 
         aligned_sequence1 = DNA('--AAAATTT')
@@ -109,10 +99,8 @@ def test_alt_gap_open_penalty(self):
         self.assertEqual(observed, expected)
 
     def test_alt_gap_extend_penalty(self):
-        s1 = DNA('AAAATTT')
-        sequence1 = DNAIterator([s1])
-        s2 = DNA('AAAAGGTTT')
-        sequence2 = DNAIterator([s2])
+        sequence1 = DNA('AAAATTT')
+        sequence2 = DNA('AAAAGGTTT')
         observed = nw_align(sequence1, sequence2, gap_open_penalty=0.01)
 
         aligned_sequence1 = DNA('AAAA-T-TT-')
@@ -121,8 +109,6 @@ def test_alt_gap_extend_penalty(self):
 
         self.assertEqual(observed, expected)
 
-        sequence1 = DNAIterator([s1])
-        sequence2 = DNAIterator([s2])
         observed = nw_align(sequence1, sequence2, gap_open_penalty=0.01,
                             gap_extend_penalty=0.001)
 
@@ -133,10 +119,8 @@ def test_alt_gap_extend_penalty(self):
         self.assertEqual(observed, expected)
 
     def test_alt_mismatch_score(self):
-        s1 = DNA('AAAATTT')
-        sequence1 = DNAIterator([s1])
-        s2 = DNA('AAAAGGTTT')
-        sequence2 = DNAIterator([s2])
+        sequence1 = DNA('AAAATTT')
+        sequence2 = DNA('AAAAGGTTT')
         observed = nw_align(sequence1, sequence2, gap_open_penalty=0.01)
 
         aligned_sequence1 = DNA('AAAA-T-TT-')
@@ -145,8 +129,6 @@ def test_alt_mismatch_score(self):
 
         self.assertEqual(observed, expected)
 
-        sequence1 = DNAIterator([s1])
-        sequence2 = DNAIterator([s2])
         observed = nw_align(sequence1, sequence2, gap_open_penalty=0.1,
                             mismatch_score=-0.1)