Add codes to members or sequences and support __radd__

BooleanCat · Jul 26, 2023 · 2825ebf · 2825ebf
1 parent 1ecb709
commit 2825ebf
Show file tree

Hide file tree

Showing 17 changed files with 341 additions and 141 deletions.
diff --git a/Makefile b/Makefile
@@ -1,6 +1,6 @@
 test: check
 	poetry run maturin develop --release
-	poetry run pytest -v .
+	poetry run pytest .
 
 check:
 	cargo fmt --check

diff --git a/haem.pyi b/haem.pyi
@@ -56,9 +56,13 @@ class DNABase:
         ...
     def __invert__(self) -> DNABase:
         """See `DNABase.complement`."""
-    def __add__(self, other: typing.Union[DNABase, DNASequence]) -> DNASequence:
-        """Create a DNA sequence from the concatenation of this base and either
-        another base or a DNA sequence."""
+    def __add__(self, other: typing.Union[DNABase, DNASequence, str]) -> DNASequence:
+        """Create a new sequence consisting of this base followed by the given
+        sequence member(s)."""
+        ...
+    def __radd__(self, other: typing.Union[DNABase, DNASequence, str]) -> DNASequence:
+        """Create a new sequence consisting of the given sequence member(s)
+        followed by this base."""
         ...
 
 class RNABase:
@@ -111,9 +115,13 @@ class RNABase:
         ...
     def __invert__(self) -> RNABase:
         """See `RNABase.complement`."""
-    def __add__(self, other: typing.Union[RNABase, RNASequence]) -> RNASequence:
-        """Create a RNA sequence from the concatenation of this base and either
-        another base or a RNA sequence."""
+    def __add__(self, other: typing.Union[RNABase, RNASequence, str]) -> RNASequence:
+        """Create a new sequence consisting of this base followed by the given
+        sequence member(s)."""
+        ...
+    def __radd__(self, other: typing.Union[RNABase, RNASequence, str]) -> RNASequence:
+        """Create a new sequence consisting of the given sequence member(s)
+        followed by this base."""
         ...
 
 class AminoAcid:
@@ -186,10 +194,16 @@ class AminoAcid:
         """Always true."""
         ...
     def __add__(
-        self, other: typing.Union[AminoAcid, AminoAcidSequence]
+        self, other: typing.Union[AminoAcid, AminoAcidSequence, str]
+    ) -> AminoAcidSequence:
+        """Create a new sequence consisting of this amino acid followed by the
+        given sequence member(s)."""
+        ...
+    def __radd__(
+        self, other: typing.Union[AminoAcid, AminoAcidSequence, str]
     ) -> AminoAcidSequence:
-        """Create an amino acid sequence from the concatenation of this amino
-        acid and either another amino acid or an amino acid sequence."""
+        """Create a new sequence consisting of the given sequence member(s)
+        followed by this amino acid."""
         ...
 
 class DNASequence:
@@ -232,9 +246,13 @@ class DNASequence:
     def __bool__(self) -> bool:
         """Casting to bool is False for empty sequences and True otherwise."""
         ...
-    def __add__(self, other: typing.Union[DNABase, DNASequence]) -> DNASequence:
-        """Create a DNA sequence from the concatenation of this sequence and
-        either a DNABase or another DNASequence."""
+    def __add__(self, other: typing.Union[DNABase, DNASequence, str]) -> DNASequence:
+        """Create a new sequence consisting of this sequence followed by the
+        given sequence member(s)."""
+    def __radd__(self, other: typing.Union[DNABase, DNASequence, str]) -> DNASequence:
+        """Create a new sequence consisting of the given sequence member(s)
+        followed by this sequence."""
+        ...
     def __contains__(self, item: typing.Union[DNABase, DNASequence]) -> bool:
         """Return true if the given DNABase or DNASequence is contained within
         this sequence."""
@@ -288,9 +306,13 @@ class RNASequence:
     def __bool__(self) -> bool:
         """Casting to bool is False for empty sequences and True otherwise."""
         ...
-    def __add__(self, other: typing.Union[RNABase, RNASequence]) -> RNASequence:
-        """Create a RNA sequence from the concatenation of this sequence and
-        either a RNABase or another RNASequence."""
+    def __add__(self, other: typing.Union[RNABase, RNASequence, str]) -> RNASequence:
+        """Create a new sequence consisting of this sequence followed by the
+        given sequence member(s)."""
+    def __radd__(self, other: typing.Union[RNABase, RNASequence, str]) -> RNASequence:
+        """Create a new sequence consisting of the given sequence member(s)
+        followed by this sequence."""
+        ...
     def __contains__(self, item: typing.Union[RNABase, RNASequence]) -> bool:
         """Return true if the given RNABase or RNASequence is contained within
         this sequence."""
@@ -337,10 +359,16 @@ class AminoAcidSequence:
         """Casting to bool is False for empty sequences and True otherwise."""
         ...
     def __add__(
-        self, other: typing.Union[AminoAcid, AminoAcidSequence]
+        self, other: typing.Union[AminoAcid, AminoAcidSequence, str]
     ) -> AminoAcidSequence:
-        """Create an amino acid sequence from the concatenation of this
-        sequence and either an AminoAcid or another AminoAcidSequence."""
+        """Create a new sequence consisting of this sequence followed by the
+        given sequence member(s)."""
+    def __radd__(
+        self, other: typing.Union[AminoAcidSequence, AminoAcid, str]
+    ) -> AminoAcidSequence:
+        """Create a new sequence consisting of the given sequence member(s)
+        followed by this sequence."""
+        ...
     def __contains__(self, item: typing.Union[AminoAcid, AminoAcidSequence]) -> bool:
         """Return true if the given AminoAcid or AminoAcidSequence is contained
         within this sequence."""

diff --git a/src/aminoacid.rs b/src/aminoacid.rs
@@ -1,6 +1,7 @@
 use crate::aminoacidsequence::AminoAcidSequence;
-use crate::member::{Member, MemberOrMembers};
+use crate::member::Member;
 use crate::rnabase::RNABase;
+use crate::utils::AddInput;
 use pyo3::class::basic::CompareOp;
 use pyo3::create_exception;
 use pyo3::prelude::*;
@@ -195,10 +196,16 @@ impl AminoAcid {
         true
     }
 
-    fn __add__(&self, other: MemberOrMembers<Self>) -> AminoAcidSequence {
-        AminoAcidSequence {
-            amino_acids: self.add(other),
-        }
+    fn __add__(&self, other: AddInput<Self>) -> PyResult<AminoAcidSequence> {
+        Ok(AminoAcidSequence {
+            amino_acids: self.add(other, false)?,
+        })
+    }
+
+    fn __radd__(&self, other: AddInput<Self>) -> PyResult<AminoAcidSequence> {
+        Ok(AminoAcidSequence {
+            amino_acids: self.add(other, true)?,
+        })
     }
 }
 

diff --git a/src/aminoacidsequence.rs b/src/aminoacidsequence.rs
@@ -1,7 +1,7 @@
 use crate::aminoacid::AminoAcid;
 use crate::member::{MemberOrCode, MemberOrMembers};
 use crate::sequence::{Sequence, SequenceInput};
-use crate::utils::IntOrSlice;
+use crate::utils::{AddInput, IntOrSlice};
 use pyo3::class::basic::CompareOp;
 use pyo3::prelude::*;
 
@@ -41,10 +41,16 @@ impl AminoAcidSequence {
         self.bool()
     }
 
-    fn __add__(&self, other: MemberOrMembers<AminoAcid>) -> Self {
-        Self {
-            amino_acids: self.add(other),
-        }
+    fn __add__(&self, other: AddInput<AminoAcid>) -> PyResult<Self> {
+        Ok(Self {
+            amino_acids: self.add(other, false)?,
+        })
+    }
+
+    fn __radd__(&self, other: AddInput<AminoAcid>) -> PyResult<Self> {
+        Ok(Self {
+            amino_acids: self.add(other, true)?,
+        })
     }
 
     fn __contains__(&self, amino_acid_or_seq: MemberOrMembers<AminoAcid>) -> PyResult<bool> {

diff --git a/src/dnabase.rs b/src/dnabase.rs
@@ -1,6 +1,7 @@
 use crate::dnasequence::DNASequence;
-use crate::member::{Member, MemberOrMembers};
+use crate::member::Member;
 use crate::rnabase::RNABase;
+use crate::utils::AddInput;
 use pyo3::class::basic::CompareOp;
 use pyo3::prelude::*;
 use std::fmt;
@@ -107,10 +108,16 @@ impl DNABase {
         self.get_complement()
     }
 
-    fn __add__(&self, other: MemberOrMembers<Self>) -> DNASequence {
-        DNASequence {
-            bases: self.add(other),
-        }
+    fn __add__(&self, other: AddInput<Self>) -> PyResult<DNASequence> {
+        Ok(DNASequence {
+            bases: self.add(other, false)?,
+        })
+    }
+
+    fn __radd__(&self, other: AddInput<Self>) -> PyResult<DNASequence> {
+        Ok(DNASequence {
+            bases: self.add(other, true)?,
+        })
     }
 
     fn __str__(&self) -> String {

diff --git a/src/dnasequence.rs b/src/dnasequence.rs
@@ -3,7 +3,7 @@ use crate::member::{MemberOrCode, MemberOrMembers};
 use crate::rnabase::RNABase;
 use crate::rnasequence::RNASequence;
 use crate::sequence::{Sequence, SequenceInput};
-use crate::utils::IntOrSlice;
+use crate::utils::{AddInput, IntOrSlice};
 use pyo3::class::basic::CompareOp;
 use pyo3::prelude::*;
 use rayon::prelude::*;
@@ -61,10 +61,16 @@ impl DNASequence {
         self.bool()
     }
 
-    fn __add__(&self, other: MemberOrMembers<DNABase>) -> Self {
-        Self {
-            bases: self.add(other),
-        }
+    fn __add__(&self, other: AddInput<DNABase>) -> PyResult<Self> {
+        Ok(Self {
+            bases: self.add(other, false)?,
+        })
+    }
+
+    fn __radd__(&self, other: AddInput<DNABase>) -> PyResult<Self> {
+        Ok(Self {
+            bases: self.add(other, true)?,
+        })
     }
 
     fn __len__(&self) -> usize {

diff --git a/src/member.rs b/src/member.rs
@@ -1,27 +1,54 @@
-use crate::utils::Wrapper;
+use crate::utils::{AddInput, Wrapper};
 use pyo3::class::basic::CompareOp;
 use pyo3::prelude::*;
 use pyo3::pyclass::PyClass;
 use pyo3::types::PyIterator;
 use rayon::prelude::*;
 
 pub trait Member<T> {
-    fn add(&self, other: MemberOrMembers<T>) -> Vec<T>;
+    fn add(&self, other: AddInput<T>, swap: bool) -> PyResult<Vec<T>>;
     fn richcmp(&self, other: &Self, op: CompareOp, py: Python<'_>) -> PyObject;
 }
 
-impl<T: Clone + PartialEq> Member<T> for T {
-    fn add(&self, other: MemberOrMembers<T>) -> Vec<T> {
-        match other {
-            MemberOrMembers::Member(member) => {
-                vec![self.clone(), member]
+impl<T: Clone + PartialEq> Member<T> for T
+where
+    T: TryFrom<char, Error = PyErr> + Send + Clone,
+{
+    fn add(&self, other: AddInput<T>, swap: bool) -> PyResult<Vec<T>> {
+        Ok(match other {
+            AddInput::Member(member) => match swap {
+                true => vec![member, self.clone()],
+                false => vec![self.clone(), member],
+            },
+            AddInput::Members(members) => {
+                let mut sequence = Vec::with_capacity(members.len() + 1);
+                match swap {
+                    false => {
+                        sequence.push(self.clone());
+                        sequence.extend(members);
+                    }
+                    true => {
+                        sequence.extend(members);
+                        sequence.push(self.clone());
+                    }
+                }
+                sequence
             }
-            MemberOrMembers::Sequence(sequence) => {
-                let mut sequence = sequence;
-                sequence.insert(0, self.clone());
+            AddInput::Codes(codes) => {
+                let mut sequence = Vec::with_capacity(codes.len() + 1);
+                match swap {
+                    false => {
+                        sequence.push(self.clone());
+                        sequence.extend(Wrapper::try_from(codes)?.into_inner());
+                    }
+                    true => {
+                        sequence.extend(Wrapper::try_from(codes)?.into_inner());
+                        sequence.push(self.clone());
+                    }
+                }
                 sequence
             }
-        }
+        })
     }
 
     fn richcmp(&self, other: &Self, op: CompareOp, py: Python<'_>) -> PyObject {
@@ -59,8 +86,8 @@ where
     fn try_from(codes: Vec<char>) -> PyResult<Self> {
         Ok(Wrapper(
             codes
-                .par_iter()
-                .map(|c| T::try_from(*c))
+                .into_par_iter()
+                .map(T::try_from)
                 .collect::<PyResult<_>>()?,
         ))
     }

diff --git a/src/rnabase.rs b/src/rnabase.rs
@@ -1,6 +1,7 @@
 use crate::dnabase::DNABase;
-use crate::member::{Member, MemberOrMembers};
+use crate::member::Member;
 use crate::rnasequence::RNASequence;
+use crate::utils::AddInput;
 use pyo3::class::basic::CompareOp;
 use pyo3::prelude::*;
 use std::fmt;
@@ -107,10 +108,16 @@ impl RNABase {
         self.get_complement()
     }
 
-    fn __add__(&self, other: MemberOrMembers<Self>) -> RNASequence {
-        RNASequence {
-            bases: self.add(other),
-        }
+    fn __add__(&self, other: AddInput<Self>) -> PyResult<RNASequence> {
+        Ok(RNASequence {
+            bases: self.add(other, false)?,
+        })
+    }
+
+    fn __radd__(&self, other: AddInput<Self>) -> PyResult<RNASequence> {
+        Ok(RNASequence {
+            bases: self.add(other, true)?,
+        })
     }
 
     fn __str__(&self) -> String {

diff --git a/src/rnasequence.rs b/src/rnasequence.rs
@@ -3,7 +3,7 @@ use crate::dnasequence::DNASequence;
 use crate::member::{MemberOrCode, MemberOrMembers};
 use crate::rnabase::RNABase;
 use crate::sequence::{Sequence, SequenceInput};
-use crate::utils::IntOrSlice;
+use crate::utils::{AddInput, IntOrSlice};
 use pyo3::class::basic::CompareOp;
 use pyo3::prelude::*;
 use rayon::prelude::*;
@@ -61,10 +61,16 @@ impl RNASequence {
         self.bool()
     }
 
-    fn __add__(&self, other: MemberOrMembers<RNABase>) -> Self {
-        Self {
-            bases: self.add(other),
-        }
+    fn __add__(&self, other: AddInput<RNABase>) -> PyResult<Self> {
+        Ok(Self {
+            bases: self.add(other, false)?,
+        })
+    }
+
+    fn __radd__(&self, other: AddInput<RNABase>) -> PyResult<Self> {
+        Ok(Self {
+            bases: self.add(other, true)?,
+        })
     }
 
     fn __len__(&self) -> usize {