Move 'export_trait_data' to 'gn3.db.traits' module

Issue:
https://github.com/genenetwork/gn-gemtext-threads/blob/main/topics/gn1-migration-to-gn2/partial-correlations.gmi

* gn3/db/traits.py: Move function `export_trait_data` here
* gn3/heatmaps.py: Remove function `export_trait_data`
* tests/unit/db/test_traits.py: Move function `export_trait_data` tests here
* tests/unit/test_heatmaps.py: Remove function `export_trait_data` here

  Function `export_trait_data` more closely corresponds to the traits and is
  used in more than just the `gn3.heatmaps` module. This commit moves the
  relevant code over to the `gn3.db.traits` module and also moves the tests to
  the corresponding tests modules.

1 files changed, 1 insertions, 66 deletions

diff --git a/gn3/heatmaps.py b/gn3/heatmaps.py
index adbfbc6..3b94e88 100644
--- a/gn3/heatmaps.py
+++ b/gn3/heatmaps.py
@@ -14,6 +14,7 @@ from plotly.subplots import make_subplots # type: ignore
 from gn3.settings import TMPDIR
 from gn3.random import random_string
 from gn3.computations.slink import slink
+from gn3.db.traits import export_trait_data
 from gn3.computations.correlations2 import compute_correlation
 from gn3.db.genotypes import (
     build_genotype_file, load_genotype_samples)
@@ -26,72 +27,6 @@ from gn3.computations.qtlreaper import (
     parse_reaper_main_results,
     organise_reaper_main_results)
 
-def export_trait_data(
-        trait_data: dict, samplelist: Sequence[str], dtype: str = "val",
-        var_exists: bool = False, n_exists: bool = False):
-    """
-    Export data according to `samplelist`. Mostly used in calculating
-    correlations.
-
-    DESCRIPTION:
-    Migrated from
-    https://github.com/genenetwork/genenetwork1/blob/master/web/webqtl/base/webqtlTrait.py#L166-L211
-
-    PARAMETERS
-    trait: (dict)
-      The dictionary of key-value pairs representing a trait
-    samplelist: (list)
-      A list of sample names
-    dtype: (str)
-      ... verify what this is ...
-    var_exists: (bool)
-      A flag indicating existence of variance
-    n_exists: (bool)
-      A flag indicating existence of ndata
-    """
-    def __export_all_types(tdata, sample):
-        sample_data = []
-        if tdata[sample]["value"]:
-            sample_data.append(tdata[sample]["value"])
-            if var_exists:
-                if tdata[sample]["variance"]:
-                    sample_data.append(tdata[sample]["variance"])
-                else:
-                    sample_data.append(None)
-            if n_exists:
-                if tdata[sample]["ndata"]:
-                    sample_data.append(tdata[sample]["ndata"])
-                else:
-                    sample_data.append(None)
-        else:
-            if var_exists and n_exists:
-                sample_data += [None, None, None]
-            elif var_exists or n_exists:
-                sample_data += [None, None]
-            else:
-                sample_data.append(None)
-
-        return tuple(sample_data)
-
-    def __exporter(accumulator, sample):
-        # pylint: disable=[R0911]
-        if sample in trait_data["data"]:
-            if dtype == "val":
-                return accumulator + (trait_data["data"][sample]["value"], )
-            if dtype == "var":
-                return accumulator + (trait_data["data"][sample]["variance"], )
-            if dtype == "N":
-                return accumulator + (trait_data["data"][sample]["ndata"], )
-            if dtype == "all":
-                return accumulator + __export_all_types(trait_data["data"], sample)
-            raise KeyError("Type `%s` is incorrect" % dtype)
-        if var_exists and n_exists:
-            return accumulator + (None, None, None)
-        if var_exists or n_exists:
-            return accumulator + (None, None)
-        return accumulator + (None,)
-
-    return reduce(__exporter, samplelist, tuple())
 
 def trait_display_name(trait: Dict):
     """