diff --git a/gnomad/resources/resource_utils.py b/gnomad/resources/resource_utils.py index f245ce70d..2c6eb64a9 100644 --- a/gnomad/resources/resource_utils.py +++ b/gnomad/resources/resource_utils.py @@ -109,16 +109,23 @@ class TableResource(BaseResource): expected_file_extensions: List[str] = [".ht"] - def ht(self, force_import: bool = False) -> hl.Table: + def ht( + self, + force_import: bool = False, + read_args: Optional[Dict[str, Any]] = None, + ) -> hl.Table: """ Read and return the Hail Table resource. + :param force_import: If ``True``, force the import of the resource even if it + already exists. + :param read_args: Any additional arguments to pass to hl.read_table. :return: Hail Table resource """ if self.path is None or force_import: return self.import_func(**self.import_args) else: - return hl.read_table(self.path) + return hl.read_table(self.path, **(read_args or {})) def import_resource(self, overwrite: bool = True, **kwargs) -> None: """ @@ -144,16 +151,23 @@ class MatrixTableResource(BaseResource): expected_file_extensions: List[str] = [".mt"] - def mt(self, force_import: bool = False) -> hl.MatrixTable: + def mt( + self, + force_import: bool = False, + read_args: Optional[Dict[str, Any]] = None, + ) -> hl.MatrixTable: """ Read and return the Hail MatrixTable resource. + :param force_import: If ``True``, force the import of the resource even if it + already exists. + :param read_args: Any additional arguments to pass to hl.read_matrix_table. :return: Hail MatrixTable resource """ if self.path is None or force_import: return self.import_func(**self.import_args) else: - return hl.read_matrix_table(self.path) + return hl.read_matrix_table(self.path, **(read_args or {})) def import_resource(self, overwrite: bool = True, **kwargs) -> None: """ @@ -179,16 +193,23 @@ class VariantDatasetResource(BaseResource): expected_file_extensions: List[str] = [".vds"] - def vds(self, force_import: bool = False) -> hl.vds.VariantDataset: + def vds( + self, + force_import: bool = False, + read_args: Optional[Dict[str, Any]] = None, + ) -> hl.vds.VariantDataset: """ Read and return the Hail VariantDataset resource. + :param force_import: If ``True``, force the import of the resource even if it + already exists. + :param read_args: Any additional arguments to pass to hl.vds.read_vds. :return: Hail VariantDataset resource """ if self.path is None or force_import: return self.import_func(**self.import_args) else: - return hl.vds.read_vds(self.path) + return hl.vds.read_vds(self.path, **(read_args or {})) def import_resource(self, overwrite: bool = True, **kwargs) -> None: """ @@ -283,13 +304,14 @@ class BlockMatrixResource(BaseResource): expected_file_extensions: List[str] = [".bm"] - def bm(self) -> BlockMatrix: + def bm(self, read_args: Optional[Dict[str, Any]] = None) -> BlockMatrix: """ Read and return the Hail MatrixTable resource. + :param read_args: Any additional arguments to pass to BlockMatrix.read. :return: Hail MatrixTable resource """ - return BlockMatrix.read(self.path) + return BlockMatrix.read(self.path, **(read_args or {})) def import_resource(self, overwrite: bool = True, **kwargs) -> None: """ @@ -318,16 +340,23 @@ class ExpressionResource(BaseResource): expected_file_extensions: List[str] = [".he"] - def he(self, force_import: bool = False) -> hl.expr.Expression: + def he( + self, + force_import: bool = False, + read_args: Optional[Dict[str, Any]] = None, + ) -> hl.expr.Expression: """ Read and return the Hail Expression resource. + :param force_import: If ``True``, force the import of the resource even if it + already exists. + :param read_args: Any additional arguments to pass to hl.experimental.read_expression. :return: Hail Expression resource. """ if self.path is None or force_import: return self.import_func(**self.import_args) else: - return hl.experimental.read_expression(self.path) + return hl.experimental.read_expression(self.path, **(read_args or {})) def import_resource(self, overwrite: bool = True, **kwargs) -> None: """