From f71b853deea708b6471a70073f128ab5a8fee1e4 Mon Sep 17 00:00:00 2001
From: Arya Massarat <23412689+aryarm@users.noreply.github.com>
Date: Sat, 9 Apr 2022 16:03:27 -0700
Subject: [PATCH] support gz and bz2 file extensions (see #19)

---
 haptools/data/covariates.py | 4 +++-
 haptools/data/phenotypes.py | 4 +++-
 2 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/haptools/data/covariates.py b/haptools/data/covariates.py
index 2ce829e3..74f78d14 100644
--- a/haptools/data/covariates.py
+++ b/haptools/data/covariates.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 from csv import reader
 from pathlib import Path
+from fileinput import hook_compressed
 
 import numpy as np
 
@@ -71,7 +72,8 @@ def read(self, samples: list[str] = None):
         """
         super().read()
         # load all info into memory
-        with open(self.fname) as covars:
+        # use hook_compressed to automatically handle gz files
+        with hook_compressed(self.fname, mode='rt') as covars:
             covar_text = reader(covars, delimiter="\t")
             header = next(covar_text)
             # there should at least two columns
diff --git a/haptools/data/phenotypes.py b/haptools/data/phenotypes.py
index 9220b90b..3c09d8a0 100644
--- a/haptools/data/phenotypes.py
+++ b/haptools/data/phenotypes.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 from csv import reader
 from pathlib import Path
+from fileinput import hook_compressed
 
 import numpy as np
 
@@ -71,7 +72,8 @@ def read(self, samples: list[str] = None):
         """
         super().read()
         # load all info into memory
-        with open(self.fname) as phens:
+        # use hook_compressed to automatically handle gz files
+        with hook_compressed(self.fname, mode='rt') as phens:
             phen_text = reader(phens, delimiter="\t")
             # convert to list and subset samples if need be
             if samples: