1 files changed, 45 insertions, 38 deletions
diff --git a/wqflask/base/species.py b/wqflask/base/species.py
index e3c29916..f303aabb 100644
--- a/wqflask/base/species.py
+++ b/wqflask/base/species.py
@@ -1,59 +1,66 @@
-import collections
-
-from flask import Flask, g
-
-
-from utility.logger import getLogger
-logger = getLogger(__name__)
+from collections import OrderedDict
+from dataclasses import dataclass
+from dataclasses import InitVar
+from typing import Optional, Dict
+from flask import g
 
 
+@dataclass
 class TheSpecies:
-    def __init__(self, dataset=None, species_name=None):
-        if species_name != None:
-            self.name = species_name
+    """Data related to species."""
+    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.dataset = dataset
             self.chromosomes = Chromosomes(dataset=self.dataset)
 
 
+@dataclass
 class IndChromosome:
-    def __init__(self, name, length):
-        self.name = name
-        self.length = length
+    """Data related to IndChromosome"""
+    name: str
+    length: int
 
     @property
     def mb_length(self):
-        """Chromosome length in megabases"""
+        """Chromosome length in mega-bases"""
         return self.length / 1000000
 
 
+@dataclass
 class Chromosomes:
-    def __init__(self, dataset=None, species=None):
-        self.chromosomes = collections.OrderedDict()
-        if species != None:
-            query = """
-                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
-                """ % species.capitalize()
-        else:
+    """Data related to a chromosome"""
+    dataset: InitVar[Dict] = None
+    species: Optional[str] = None
+
+    def __post_init__(self, dataset):
+        if self.species is None:
             self.dataset = dataset
 
-            query = """
-                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
-        logger.sql(query)
+    @property
+    def chromosomes(self):
+        """Lazily fetch the chromosomes"""
+        chromosomes = OrderedDict()
+        if self.species is not None:
+            query = (
+                "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:
+            query = (
+                "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 = g.db.execute(query).fetchall()
-
         for item in results:
-            self.chromosomes[item.OrderId] = IndChromosome(
+            chromosomes[item.OrderId] = IndChromosome(
                 item.Name, item.Length)
+        return chromosomes