Integrate get_lsr_from_chr function

* gn3/heatmaps.py: copy over function
* tests/unit/test_heatmaps.py: add tests

  Copy function over from proof of concept and add some tests to ensure it
  works as expected.

2 files changed, 22 insertions, 0 deletions

diff --git a/gn3/heatmaps.py b/gn3/heatmaps.py
index 198fb45..991ddec 100644
--- a/gn3/heatmaps.py
+++ b/gn3/heatmaps.py
@@ -276,6 +276,14 @@ def get_nearest_marker(traits_list, genotype):
     marker_finder = nearest_marker_finder(genotype)
     return [marker_finder(trait) for trait in traits_list]
 
+def get_lrs_from_chr(trait, chr_name):
+    chromosome = trait["chromosomes"].get(chr_name)
+    if chromosome:
+        return [
+            locus["LRS"] for locus in
+            sorted(chromosome["loci"], key=lambda loc: loc["Locus"])]
+    return [None]
+
 # # Grey + Blue + Red
 # def generate_heatmap():
 #     cols = 20
diff --git a/tests/unit/test_heatmaps.py b/tests/unit/test_heatmaps.py
index 265d5a8..cfdde1e 100644
--- a/tests/unit/test_heatmaps.py
+++ b/tests/unit/test_heatmaps.py
@@ -2,6 +2,7 @@
 from unittest import TestCase
 from gn3.heatmaps import (
     cluster_traits,
+    get_lrs_from_chr,
     export_trait_data,
     compute_traits_order,
     retrieve_strains_and_values)
@@ -185,3 +186,16 @@ class TestHeatmap(TestCase):
             with self.subTest(strainlist=slist, traitdata=tdata):
                 self.assertEqual(
                     retrieve_strains_and_values(orders, slist, tdata), expected)
+
+    def test_get_lrs_from_chr(self):
+        for trait, chromosome, expected in [
+                [{"chromosomes": {}}, 3, [None]],
+                [{"chromosomes": {3: {"loci": [
+                    {"Locus": "b", "LRS": 1.9},
+                    {"Locus": "a", "LRS": 13.2},
+                    {"Locus": "d", "LRS": 53.21},
+                    {"Locus": "c", "LRS": 2.22}]}}},
+                 3,
+                 [13.2, 1.9, 2.22, 53.21]]]:
+            with self.subTest(trait=trait, chromosome=chromosome):
+                self.assertEqual(get_lrs_from_chr(trait, chromosome), expected)