Generalise fetching of samples/cases/individuals.

1 files changed, 17 insertions, 9 deletions

diff --git a/r_qtl/r_qtl2.py b/r_qtl/r_qtl2.py
index f03aff5..1e28bc0 100644
--- a/r_qtl/r_qtl2.py
+++ b/r_qtl/r_qtl2.py
@@ -368,19 +368,27 @@ def read_geno_file_data(
             replace_genotype_codes, genocodes=cdata.get("genotypes", {})))
 
 
-def load_geno_samples(zipfilepath: Union[str, Path]) -> tuple[str, ...]:
-    """Load the samples/cases/individuals from the 'geno' file(s)."""
+def load_samples(
+        zipfilepath: Union[str, Path], filetype: str) -> tuple[str, ...]:
+    """Load the samples/cases/individuals from file(s) of type 'filetype'."""
     cdata = read_control_file(zipfilepath)
-    samples = set()
-    for genofile in cdata.get("geno", []):
-        gdata = read_geno_file_data(zipfilepath, genofile)
-        if cdata.get("geno_transposed", False):
-            samples.update(next(gdata)[1:])
+    samples: set[str] = set()
+    for afile in cdata.get(filetype, []):
+        filedata = read_geno_file_data(zipfilepath, afile)
+        if cdata.get(f"{filetype}_transposed", False):
+            samples.update(
+                item for item in next(filedata)[1:] if item is not None)
         else:
             try:
-                next(gdata)# Ignore first row.
-                samples.update(line[0] for line in gdata)
+                next(filedata)# Ignore first row.
+                samples.update(
+                    line[0] for line in filedata if line[0] is not None)
             except StopIteration:# Empty file.
                 pass
 
     return tuple(samples)
+
+
+load_geno_samples = partial(load_samples, filetype="geno")
+load_founder_geno_samples = partial(load_samples, filetype="founder_geno")
+load_pheno_samples = partial(load_samples, filetype="pheno")