Repository: cogent3/cogent3
Pull request #1662 (2023-12-09 05:55:57)

ENH: enable indexing seqs with numpy int dtypes


Patch:
diff --git a/src/cogent3/core/sequence.py b/src/cogent3/core/sequence.py
index 5a78c9489..690f3cffb 100644
--- a/src/cogent3/core/sequence.py
+++ b/src/cogent3/core/sequence.py
@@ -29,6 +29,9 @@ from numpy import (
     arange,
     array,
     compress,
+    floating,
+    integer,
+    issubdtype,
     logical_not,
     logical_or,
     nonzero,
@@ -63,11 +66,7 @@ from cogent3.util.misc import (
     get_setting_from_environ,
 )
 from cogent3.util.transform import for_seq, per_shortest
-from cogent3.util.warning import (
-    deprecated,
-    deprecated_args,
-    deprecated_callable,
-)
+from cogent3.util.warning import deprecated_args
 
 
 ARRAY_TYPE = type(array(1))
@@ -77,6 +76,16 @@ frac_same = for_seq(f=eq, aggregator=sum, normalizer=per_shortest)
 frac_diff = for_seq(f=ne, aggregator=sum, normalizer=per_shortest)
 
 
+def _is_int(val) -> bool:
+    """whether val is builtin, or numpy, integer"""
+    return issubdtype(val.__class__, integer) or isinstance(val, int)
+
+
+def _is_float(val) -> bool:
+    """whether val is builtin, or numpy, integer"""
+    return issubdtype(val.__class__, floating) or isinstance(val, float)
+
+
 @total_ordering
 class SequenceI(object):
     """Abstract class containing Sequence interface.
@@ -1317,7 +1326,7 @@ class Sequence(SequenceI):
             new = self._mapped(index)
             preserve_offset = not index.reverse
 
-        elif isinstance(index, (int, slice)):
+        elif isinstance(index, slice) or _is_int(index):
             new = self.__class__(
                 self._seq[index], name=self.name, check=False, info=self.info
             )
@@ -1331,6 +1340,9 @@ class Sequence(SequenceI):
             new.replace_annotation_db(self.annotation_db, check=False)
             new.annotation_offset = self.annotation_offset
 
+        if _is_float(index):
+            raise TypeError("cannot slice using float")
+
         if hasattr(self, "_repr_policy"):
             new._repr_policy.update(self._repr_policy)
 
@@ -2139,7 +2151,7 @@ class SeqView:
         )
 
     def __getitem__(self, segment):
-        if isinstance(segment, int):
+        if _is_int(segment):
             start, stop, step = self._get_index(segment)
             return self.__class__(self.seq, start=start, stop=stop, step=step)
 


Test patch:
diff --git a/tests/test_core/test_sequence.py b/tests/test_core/test_sequence.py
index f49a68e7f..084bed310 100644
--- a/tests/test_core/test_sequence.py
+++ b/tests/test_core/test_sequence.py
@@ -7,6 +7,7 @@ import re
 from pickle import dumps
 from unittest import TestCase
 
+import numpy
 import pytest
 
 from numpy import array
@@ -2461,3 +2462,18 @@ def test_trim_terminal_stop_strict():
     seq = cogent3.make_seq("TCCAG", moltype="dna")
     with pytest.raises(AlphabetError):
         seq.trim_stop_codon(gc=gc, strict=True)
+
+
+@pytest.mark.parametrize("cast", (int, numpy.int32, numpy.int64, numpy.uint8))
+def test_index_a_seq(cast):
+    seq = cogent3.make_seq("TCCAG", moltype="dna")
+    got = seq[cast(1)]
+    assert isinstance(got, Sequence)
+
+
+@pytest.mark.parametrize("cast", (float, numpy.float32))
+def test_index_a_seq_float_fail(cast):
+    seq = cogent3.make_seq("TCCAG", moltype="dna")
+    index = cast(1)
+    with pytest.raises(TypeError):
+        seq[index]