added tests

Author: Avsecz

kipoiseq/dataclasses.py

@@ -42,10 +42,10 @@ def __init__(self,
         # main 4 attributes
         # these should be immutable by default to not
         # run into any strange issues downstream.
-        self._chrom = chrom
-        self._pos = pos
-        self._ref = ref
-        self._alt = alt
+        self._chrom = str(chrom)
+        self._pos = int(pos)
+        self._ref = str(ref)
+        self._alt = str(alt)
 
         # other 4 main VCF attributes
         self.id = id
@@ -109,15 +109,16 @@ def __hash__(self):
         return hash((self.chrom, self.pos, self.ref, self.alt))
 
     def __str__(self):
-        return f"{self.chrom}_{self.pos}_{self.ref}/{self.alt}"
+        return f"{self.chrom}:{self.pos}:{self.ref}>{self.alt}"
 
-    def from_str(self, s):
-        chrom, pos, ref_alt = s.split("_")
-        ref, alt = ref_alt.split("/")
-        return Variant(chrom=chrom, pos=pos, ref=ref, alt=alt)
+    @classmethod
+    def from_str(cls, s):
+        chrom, pos, ref_alt = s.split(":")
+        ref, alt = ref_alt.split(">")
+        return cls(chrom=chrom, pos=int(pos), ref=ref, alt=alt)
 
     def __repr__(self):
-        return f"Variant(chrom='{self.variant}'), pos={self.pos}, ref='{self.ref}', alt='{self.alt}', id='{self.id}',..."
+        return f"Variant(chrom='{self.chrom}', pos={self.pos}, ref='{self.ref}', alt='{self.alt}', id='{self.id}',...)"
 
 
 class Interval:
@@ -206,18 +207,18 @@ def to_pybedtools(self):
     def neg_strand(self):
         return self.strand == "-"
 
-    def center(self, ignore_strand=False):
+    def center(self, use_strand=True):
         """Compute the center of the interval
         """
-        if ignore_strand:
-            add_offset = 0
-        else:
+        if use_strand:
             add_offset = 0 if self.neg_strand else 1
+        else:
+            add_offset = 0
         delta = (self.end + self.start) % 2
         center = (self.end + self.start) // 2
         return center + add_offset * delta
 
-    def shift(self, x: int, use_strand: bool=False):
+    def shift(self, x: int, use_strand: bool=True):
         """Shift the interval by x.
 
         Args:
@@ -241,9 +242,9 @@ def shift(self, x: int, use_strand: bool=False):
     def swap_strand(self):
         obj = self.copy()
         if obj.strand == "+":
-            obj.strand = "-"
+            obj._strand = "-"
         elif obj.strand == "-":
-            obj.strand = "+"
+            obj._strand = "+"
         return obj
 
     def __eq__(self, obj):
@@ -261,18 +262,19 @@ def __str__(self):
     def __repr__(self):
         return (f"Interval(chrom='{self.chrom}', start={self.start}, end={self.end}, name='{self.name}', strand='{self.strand}', ...)")
 
-    def from_str(self, s):
+    @classmethod
+    def from_str(cls, s):
         chrom, int_range, strand = s.split(":")
         start, end = int_range.split("-")
-        return Interval(chrom=chrom,
-                        start=int(start),
-                        end=int(end),
-                        strand=strand)
+        return cls(chrom=chrom,
+                   start=int(start),
+                   end=int(end),
+                   strand=strand)
 
     def copy(self):
         return deepcopy(self)
 
-    def slop(self, upstream=0, downstream=0, use_strand=False):
+    def slop(self, upstream=0, downstream=0, use_strand=True):
         """Extend the interval on each strand
         """
         obj = self.copy()
@@ -296,20 +298,21 @@ def truncate(self, chrom_len=math.inf):
             obj._end = min(self.end, chrom_len - 1)
             return obj
 
-    def resize(self, width):
+    def resize(self, width, use_strand=True):
         obj = deepcopy(self)
 
         if width is None or self.width() == width:
             # no need to resize
             return obj
 
-        if self.strand != "-":
+        if self.neg_strand and use_strand:
+            # negative strand
+            obj._start = self.center() - width // 2
+            obj._end = self.center() + width // 2 + width % 2
+        else:
             # positive strand
             obj._start = self.center() - width // 2 - width % 2
             obj._end = self.center() + width // 2
-        else:
-            obj._start = self.center() - width // 2
-            obj._end = self.center() + width // 2 + width % 2
         return obj
 
     def width(self):
@@ -318,12 +321,12 @@ def width(self):
     def __len__(self):
         return self.width()
 
-    def trim(self, i, j):
+    def trim(self, i, j, use_strand=True):
         if i == 0 and j == self.width():
             return self
         obj = self.copy()
         assert j > i
-        if self.strand == "-":
+        if self.strand == "-" and use_strand:
             w = self.width()
             obj._start = self.start + w - j
             obj._end = self.start + w - i

tests/test_dataclasses.py

@@ -8,20 +8,179 @@
 - =, hashable
 -
 # VCF
-- from cyvcf2 
+- from cyvcf2
 
 # Interval
 - validation of the interval
 - access to all the attributes
 - from_pybedtools and to_pybedtools
 - shift, swapt_strand, trim, etc
 """
-
-
 from kipoiseq.dataclasses import Variant, Interval
+import pytest
+import cyvcf2
+import pybedtools
+
 
 def test_variant():
     v = Variant("chr1", 10, 'C', 'T')
 
-    #
-    pass
+    assert v.start == 9
+    assert v.chrom == 'chr1'
+    assert v.pos == 10
+    assert v.ref == 'C'
+    assert v.alt == 'T'
+    assert isinstance(v.info, dict)
+    assert len(v.info) == 0
+    assert v.qual == 0
+    assert v.filter == 'PASS'
+    v.info['test'] = 10
+    assert v.info['test'] == 10
+    assert isinstance(str(v), str)
+
+    # make sure the original got unchangd
+    v2 = v.copy()
+    v.info['test'] = 20
+    assert v2.info['test'] == 10
+    v.__repr__()
+
+    # __str__, from_str
+    assert v == Variant.from_str(str(v))
+
+    # hash test
+    assert isinstance(hash(v), int)
+    assert hash(v) == hash(Variant.from_str(str(v)))
+
+    # fixed arguments
+    with pytest.raises(AttributeError):
+        v.chrom = 'asd'
+    with pytest.raises(AttributeError):
+        v.pos = 10
+    with pytest.raises(AttributeError):
+        v.ref = 'asd'
+    with pytest.raises(AttributeError):
+        v.alt = 'asd'
+
+    # non-fixed arguments
+    v.id = 'asd'
+    v.qual = 10
+    v.filter = 'asd'
+    v.source = 2
+
+    assert isinstance(Variant("chr1", '10', 'C', 'T').pos, int)
+
+    # from cyvcf2
+    vcf = cyvcf2.VCF('tests/data/test.vcf.gz')
+    cv = list(vcf)[0]
+
+    v2 = Variant.from_cyvcf(cv)
+    assert isinstance(v2.source, cyvcf2.Variant)
+
+
+def test_interval():
+    interval = Interval("chr1", 10, 20, strand='-')
+    interval.__repr__()
+
+    assert interval.start == 10
+    assert interval.end == 20
+    assert interval.chrom == 'chr1'
+    assert interval.name == ''
+    assert isinstance(interval.attrs, dict)
+    assert len(interval.attrs) == 0
+    interval.attrs['test'] = 10
+    assert interval.attrs['test'] == 10
+    assert isinstance(str(interval), str)
+    assert interval.neg_strand
+
+    assert interval.width() == 10
+    assert len(interval) == 10
+
+    # __str__, from_str
+    assert interval == Interval.from_str(str(interval))
+
+    # make sure the original got unchangd
+    i2 = interval.copy()
+    interval.attrs['test'] = 20
+    assert i2.attrs['test'] == 10
+
+    # hash test
+    assert isinstance(hash(interval), int)
+    assert hash(interval) == hash(Interval.from_str(str(interval)))
+
+    # fixed arguments
+    with pytest.raises(AttributeError):
+        interval.chrom = 'asd'
+    with pytest.raises(AttributeError):
+        interval.start = 10
+    with pytest.raises(AttributeError):
+        interval.end = 300
+    with pytest.raises(AttributeError):
+        interval.strand = '+'
+    assert interval.strand == '-'
+
+    # non-fixed arguments
+    interval.name = 'asd'
+    interval.score = 10
+
+    assert interval == Interval.from_pybedtools(interval.to_pybedtools())
+    assert isinstance(interval.to_pybedtools(), pybedtools.Interval)
+
+    i2 = interval.shift(10, use_strand=False)
+
+    # original unchanged
+    assert interval.start == 10
+    assert interval.end == 20
+
+    assert i2.start == 20
+    assert i2.end == 30
+
+    i2 = interval.shift(10)  # use_strand = True by default
+    assert i2.start == 0
+    assert i2.end == 10
+
+    assert not interval.shift(20, use_strand=True).is_valid()
+
+    i2 = interval.shift(15, use_strand=True).truncate()
+    assert i2.start == 0
+    assert i2.end == 5
+
+    assert interval.center() == 15
+
+    # resize
+    i2 = interval.resize(11)
+    assert i2.start == 10 and i2.end == 21
+
+    i2 = interval.resize(12)
+    assert i2.start == 9 and i2.end == 21
+
+    i2 = interval.resize(9)
+    assert i2.start == 11 and i2.end == 20
+
+    i2 = interval.swap_strand()
+    assert interval.strand == "-"
+    assert i2.strand == "+"
+    assert i2.strand == '+'
+    assert len(i2) == 10
+    i2 = i2.resize(11)
+    assert i2.start == 9 and i2.end == 20
+
+    i2 = interval.copy()
+    assert i2.center(use_strand=True) == 15
+    assert i2.center(use_strand=False) == 15
+
+    i2 = interval.swap_strand()
+    assert i2.strand == "+"
+
+    i3 = i2.resize(11).shift(1)
+    assert i3.center(use_strand=True) == 16
+    assert i3.center(use_strand=False) == 15
+
+    i3 = i2.resize(11).shift(1)
+    assert i3.center(use_strand=True) == 16
+    assert i3.center(use_strand=False) == 15
+
+    i2 = interval.trim(1, 10)
+    assert i2.start == 10 and i2.end == 19
+
+    i2 = interval.trim(1, 10, use_strand=False)
+    assert i2.start == 11 and i2.end == 20