1 files changed, 31 insertions, 29 deletions
diff --git a/wqflask/base/species.py b/wqflask/base/species.py
index 0a13315c..403e0bbb 100644
--- a/wqflask/base/species.py
+++ b/wqflask/base/species.py
@@ -1,7 +1,7 @@
 from collections import OrderedDict
 from dataclasses import dataclass
 from dataclasses import InitVar
-from typing import Optional, Dict
+from typing import Optional, Dict, Any, Union
 from wqflask.database import database_connection
 
 
@@ -11,12 +11,13 @@ class TheSpecies:
     dataset: Optional[Dict] = None
     species_name: Optional[str] = None
 
-    def __post_init__(self):
-        if self.species_name is not None:
-            self.name = self.species_name
-            self.chromosomes = Chromosomes(species=self.name)
-        else:
-            self.chromosomes = Chromosomes(dataset=self.dataset)
+    def __post_init__(self) -> None:
+        # Just an alias of species_name.  It's safe for this to be None.
+        self.name = self.species_name
+        with database_connection() as conn:
+            self.chromosomes = Chromosomes(conn=conn,
+                                           species=self.species_name,
+                                           dataset=self.dataset)
 
 
 @dataclass
@@ -26,7 +27,7 @@ class IndChromosome:
     length: int
 
     @property
-    def mb_length(self):
+    def mb_length(self) -> Union[int, float]:
         """Chromosome length in mega-bases"""
         return self.length / 1000000
 
@@ -34,33 +35,34 @@ class IndChromosome:
 @dataclass
 class Chromosomes:
     """Data related to a chromosome"""
+    conn: Any
     dataset: InitVar[Dict] = None
     species: Optional[str] = None
 
-    def __post_init__(self, dataset):
+    def __post_init__(self, dataset) -> None:
         if self.species is None:
             self.dataset = dataset
 
     @property
-    def chromosomes(self):
+    def chromosomes(self) -> OrderedDict:
         """Lazily fetch the chromosomes"""
         chromosomes = OrderedDict()
-        if self.species is not None:
-            cursor.execute(
-                "SELECT Chr_Length.Name, Chr_Length.OrderId, Length "
-                "FROM Chr_Length, Species WHERE "
-                "Chr_Length.SpeciesId = Species.SpeciesId AND "
-                "Species.Name = %s "
-                "ORDER BY OrderId", (self.species.capitalize(),))
-        else:
-            cursor.execute(
-                "SELECT Chr_Length.Name, Chr_Length.OrderId, "
-                "Length FROM Chr_Length, InbredSet WHERE "
-                "Chr_Length.SpeciesId = InbredSet.SpeciesId AND "
-                "InbredSet.Name = "
-                "%s ORDER BY OrderId", (self.dataset.group.name,))
-        results = cursor.fetchall()
-        for item in results:
-            chromosomes[item.OrderId] = IndChromosome(
-                item.Name, item.Length)
-        return chromosomes
+        with self.conn.cursor() as cursor:
+            if self.species is not None:
+                cursor.execute(
+                    "SELECT Chr_Length.Name, Chr_Length.OrderId, Length "
+                    "FROM Chr_Length, Species WHERE "
+                    "Chr_Length.SpeciesId = Species.SpeciesId AND "
+                    "Species.Name = %s "
+                    "ORDER BY OrderId", (self.species.capitalize(),))
+            else:
+                cursor.execute(
+                    "SELECT Chr_Length.Name, Chr_Length.OrderId, "
+                    "Length FROM Chr_Length, InbredSet WHERE "
+                    "Chr_Length.SpeciesId = InbredSet.SpeciesId AND "
+                    "InbredSet.Name = "
+                    "%s ORDER BY OrderId", (self.dataset.group.name,))
+            for name, _, length in cursor.fetchall():
+                chromosomes[name] = IndChromosome(
+                    name=name, length=length)
+            return chromosomes