From f063c4092fe76b62f8f9d96a8056453713c7b40a Mon Sep 17 00:00:00 2001
From: Manuel Lera-Ramirez <manulera14@gmail.com>
Date: Tue, 8 Oct 2024 11:55:48 +0100
Subject: [PATCH] tests for #281 #279, add type to pcr (#298)

---
 src/pydna/amplify.py         |  2 +-
 src/pydna/utils.py           |  2 ++
 tests/test_module_amplify.py | 22 ++++++++++++++++++++--
 tests/test_module_utils.py   | 30 +++++++++++++++++++++++++++---
 4 files changed, 50 insertions(+), 6 deletions(-)

diff --git a/src/pydna/amplify.py b/src/pydna/amplify.py
index 3c4ebea0..3c12854a 100644
--- a/src/pydna/amplify.py
+++ b/src/pydna/amplify.py
@@ -431,7 +431,7 @@ def __str__(self):
     report = __str__
 
 
-def pcr(*args, **kwargs):
+def pcr(*args, **kwargs) -> _Amplicon:
     """pcr is a convenience function for the Anneal class to simplify its
     usage, especially from the command line. If more than one or no PCR
     product is formed, a ValueError is raised.
diff --git a/src/pydna/utils.py b/src/pydna/utils.py
index acb21c0f..b42d4696 100644
--- a/src/pydna/utils.py
+++ b/src/pydna/utils.py
@@ -81,6 +81,8 @@ def shift_location(original_location, shift, lim):
     newparts = []
     strand = original_location.strand
     if lim is None:
+        if min(original_location) + shift < 0:
+            raise ValueError("Shift moves location below zero, use a `lim` to loop around if sequence is circular.")
         lim = _sys.maxsize
 
     for part in original_location.parts:
diff --git a/tests/test_module_amplify.py b/tests/test_module_amplify.py
index 4a30c9a9..eeb4d895 100644
--- a/tests/test_module_amplify.py
+++ b/tests/test_module_amplify.py
@@ -614,13 +614,12 @@ def test_pcr():
     )
 
     for key, tst in enumerate(raw):
-        #print(tst[0], pcr(tst[1:]).seguid(), tst[0] in pcr(tst[1:]).seguid())
+        # print(tst[0], pcr(tst[1:]).seguid(), tst[0] in pcr(tst[1:]).seguid())
         assert tst[0] in pcr(tst[1:]).seguid()
 
 
 def test_shifts():
     from pydna.parsers import parse
-    from pydna.parsers import parse_primers
     from pydna.amplify import pcr
 
     # from pydna.amplify import nopcr
@@ -782,5 +781,24 @@ def test_shifts():
     f = pcr(f, r, t)
 
 
+def test_annotation():
+    """
+    Test that annotations are correctly added to the amplicon in primers with tails
+    https://github.com/BjornFJohansson/pydna/issues/279
+    """
+    from pydna.amplify import pcr
+    from pydna.dseqrecord import Dseqrecord
+
+    dsr = Dseqrecord("ATGCAAACAGTAATGATGGATGACATTCAAAGCACTGATTCTATTGCTGAAAAAGATAAT")
+    dsr.add_feature(x=0, y=60, type="gene", label="my_gene")  # We add a feature to highlight the sequence as a gene
+
+    forward_primer = "ccccggatccATGCAAACAGTAATGATGGA"
+    reverse_primer = "ttttggatccATTATCTTTTTCAGCAATAGAATCA"
+
+    pcr_product = pcr(forward_primer, reverse_primer, dsr)
+
+    assert pcr_product.features[0].location.extract(pcr_product).seq == dsr.seq
+
+
 if __name__ == "__main__":
     pytest.main([__file__, "-vv", "-s"])
diff --git a/tests/test_module_utils.py b/tests/test_module_utils.py
index d733720d..6a25d74f 100644
--- a/tests/test_module_utils.py
+++ b/tests/test_module_utils.py
@@ -407,7 +407,6 @@ def test_smallest_rotation():
 
 
 def test_memorize(monkeypatch):
-    import pytest
     from unittest import mock
 
     from pydna.utils import memorize as _memorize
@@ -474,6 +473,33 @@ def test_shift_location():
         loc = SimpleLocation(0, 2, strand)
         assert shift_location(shift_location(loc, 1, 6), -1, 6) == loc
 
+    # Shifting location on circular sequence
+    for strand in (1, -1, None):
+        loc = SimpleLocation(0, 4, strand)
+        assert shift_location(loc, 1, 6) == SimpleLocation(1, 5, strand)
+        if strand == -1:
+            assert shift_location(loc, -1, 6) == SimpleLocation(0, 3, strand) + SimpleLocation(5, 6, strand)
+        else:
+            assert shift_location(loc, -1, 6) == SimpleLocation(5, 6, strand) + SimpleLocation(0, 3, strand)
+
+    # Shifting ignoring the sequence length
+    # See https://github.com/BjornFJohansson/pydna/issues/281
+    for strand in (1, -1, None):
+        loc = SimpleLocation(4, 6, strand)
+        assert shift_location(loc, 1000, None) == SimpleLocation(1004, 1006, strand)
+        assert shift_location(loc, -4, None) == SimpleLocation(0, 2, strand)
+        try:
+            shift_location(loc, -1000, None)
+            raise AssertionError("Shift below zero should raise ValueError")
+        except ValueError:
+            pass
+
+        composed_loc = SimpleLocation(2, 4, strand) + SimpleLocation(5, 6, strand)
+        assert shift_location(composed_loc, 1000, None) == SimpleLocation(1002, 1004, strand) + SimpleLocation(
+            1005, 1006, strand
+        )
+        assert shift_location(composed_loc, -2, None) == SimpleLocation(0, 2, strand) + SimpleLocation(3, 4, strand)
+
 
 def test_locations_overlap():
     from pydna.utils import locations_overlap, shift_location
@@ -509,7 +535,5 @@ def test_locations_overlap():
             assert not locations_overlap(main_shifted, loc_shifted, 20)
 
 
-
-
 if __name__ == "__main__":
     pytest.main([__file__, "-vv", "-s"])