Merge pull request #219 from jkmckenna/codex/refactor-branch-for-optional-dependencies

Standardize optional dependency loading with `require`, add TYPE_CHECKING hints, and simplify lazy exports

Author: jkmckenna

src/smftools/cli/hmm_adata.py

@@ -3,18 +3,22 @@
 import copy
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Any, List, Optional, Sequence, Tuple, Union
+from typing import Any, List, Optional, Sequence, Tuple, Union, TYPE_CHECKING
 
 import numpy as np
-import torch
-
 from smftools.logging_utils import get_logger
+from smftools.optional_imports import require
 
 # FIX: import _to_dense_np to avoid NameError
 from ..hmm.HMM import _safe_int_coords, _to_dense_np, create_hmm, normalize_hmm_feature_sets
 
 logger = get_logger(__name__)
 
+if TYPE_CHECKING:
+    import torch as torch_types
+
+torch = require("torch", extra="torch", purpose="HMM CLI")
+
 # =============================================================================
 # Helpers: extracting training arrays
 # =============================================================================

src/smftools/cli/spatial_adata.py

@@ -6,6 +6,7 @@
 import anndata as ad
 
 from smftools.logging_utils import get_logger
+from smftools.optional_imports import require
 
 logger = get_logger(__name__)
 
@@ -155,7 +156,8 @@ def spatial_adata_core(
 
     import numpy as np
     import pandas as pd
-    import scanpy as sc
+
+    sc = require("scanpy", extra="scanpy", purpose="spatial analyses")
 
     from ..metadata import record_smftools_metadata
     from ..plotting import (

src/smftools/hmm/HMM.py

@@ -3,14 +3,20 @@
 import ast
 import json
 from pathlib import Path
-from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
+from typing import Any, Dict, List, Optional, Sequence, Tuple, Union, TYPE_CHECKING
 
 import numpy as np
-import torch
-import torch.nn as nn
 from scipy.sparse import issparse
 
 from smftools.logging_utils import get_logger
+from smftools.optional_imports import require
+
+if TYPE_CHECKING:
+    import torch as torch_types
+    import torch.nn as nn_types
+
+torch = require("torch", extra="torch", purpose="HMM modeling")
+nn = torch.nn
 
 logger = get_logger(__name__)
 # =============================================================================

src/smftools/hmm/__init__.py

@@ -21,11 +21,4 @@ def __getattr__(name: str):
     raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
 
 
-__all__ = [
-    "call_hmm_peaks",
-    "display_hmm",
-    "load_hmm",
-    "refine_nucleosome_calls",
-    "infer_nucleosomes_in_large_bound",
-    "save_hmm",
-]
+__all__ = list(_LAZY_ATTRS.keys())

src/smftools/hmm/call_hmm_peaks.py

@@ -5,6 +5,7 @@
 from typing import Any, Dict, Optional, Sequence, Union
 
 from smftools.logging_utils import get_logger
+from smftools.optional_imports import require
 
 logger = get_logger(__name__)
 
@@ -36,12 +37,13 @@ def call_hmm_peaks(
       - adata.var["is_in_any_{layer}_peak_{ref}"]
       - adata.var["is_in_any_peak"] (global)
     """
-    import matplotlib.pyplot as plt
     import numpy as np
     import pandas as pd
     from scipy.signal import find_peaks
     from scipy.sparse import issparse
 
+    plt = require("matplotlib.pyplot", extra="plotting", purpose="HMM peak plots")
+
     if not inplace:
         adata = adata.copy()
 

src/smftools/hmm/display_hmm.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from smftools.logging_utils import get_logger
+from smftools.optional_imports import require
 
 logger = get_logger(__name__)
 
@@ -13,7 +14,7 @@ def display_hmm(hmm, state_labels=["Non-Methylated", "Methylated"], obs_labels=[
         state_labels: Optional labels for states.
         obs_labels: Optional labels for observations.
     """
-    import torch
+    torch = require("torch", extra="torch", purpose="HMM display")
 
     logger.info("**HMM Model Overview**")
     logger.info("%s", hmm)

src/smftools/hmm/hmm_readwrite.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+from smftools.optional_imports import require
+
 
 def load_hmm(model_path, device="cpu"):
     """
@@ -8,7 +10,7 @@ def load_hmm(model_path, device="cpu"):
     Parameters:
         model_path (str): Path to a pretrained HMM
     """
-    import torch
+    torch = require("torch", extra="torch", purpose="HMM read/write")
 
     # Load model using PyTorch
     hmm = torch.load(model_path)
@@ -23,6 +25,6 @@ def save_hmm(model, model_path):
         model: HMM model instance.
         model_path: Output path for the model.
     """
-    import torch
+    torch = require("torch", extra="torch", purpose="HMM read/write")
 
     torch.save(model, model_path)

src/smftools/informatics/converted_BAM_to_adata.py

@@ -7,14 +7,13 @@
 import traceback
 from multiprocessing import Manager, Pool, current_process
 from pathlib import Path
-from typing import Iterable, Optional, Union
+from typing import Iterable, Optional, Union, TYPE_CHECKING
 
 import anndata as ad
 import numpy as np
 import pandas as pd
-import torch
-
 from smftools.logging_utils import get_logger
+from smftools.optional_imports import require
 
 from ..readwrite import make_dirs
 from .bam_functions import count_aligned_reads, extract_base_identities
@@ -24,6 +23,11 @@
 
 logger = get_logger(__name__)
 
+if TYPE_CHECKING:
+    import torch as torch_types
+
+torch = require("torch", extra="torch", purpose="converted BAM processing")
+
 if __name__ == "__main__":
     multiprocessing.set_start_method("forkserver", force=True)
 

src/smftools/informatics/fasta_functions.py

@@ -1,39 +1,44 @@
 from __future__ import annotations
 
 import gzip
+import shutil
+import subprocess
 from concurrent.futures import ProcessPoolExecutor
+from importlib.util import find_spec
 from pathlib import Path
-from typing import Dict, Iterable, Tuple
+from typing import Dict, Iterable, Tuple, TYPE_CHECKING
 
 import numpy as np
 from Bio import SeqIO
 from Bio.Seq import Seq
 from Bio.SeqRecord import SeqRecord
 
 from smftools.logging_utils import get_logger
+from smftools.optional_imports import require
 
 from ..readwrite import time_string
 
 logger = get_logger(__name__)
 
-try:
-    import pysam
-except Exception:
-    pysam = None  # type: ignore
+if TYPE_CHECKING:
+    import pysam as pysam_module
 
-try:
-    import shutil
-    import subprocess
-except Exception:  # pragma: no cover - stdlib
-    shutil = None  # type: ignore
-    subprocess = None  # type: ignore
+
+def _require_pysam() -> "pysam_module":
+    if pysam_types is not None:
+        return pysam_types
+    return require("pysam", extra="pysam", purpose="FASTA access")
+
+pysam_types = None
+if find_spec("pysam") is not None:
+    pysam_types = require("pysam", extra="pysam", purpose="FASTA access")
 
 
 def _resolve_fasta_backend() -> str:
     """Resolve the backend to use for FASTA access."""
     if shutil is not None and shutil.which("samtools"):
         return "cli"
-    if pysam is not None:
+    if pysam_types is not None:
         return "python"
     raise RuntimeError("FASTA access requires pysam or samtools in PATH.")
 
@@ -43,10 +48,9 @@ def _ensure_fasta_index(fasta: Path) -> None:
     if fai.exists():
         return
     if subprocess is None or shutil is None or not shutil.which("samtools"):
-        if pysam is not None:
-            pysam.faidx(str(fasta))
-            return
-        raise RuntimeError("FASTA indexing requires pysam or samtools in PATH.")
+        pysam_mod = _require_pysam()
+        pysam_mod.faidx(str(fasta))
+        return
     cp = subprocess.run(
         ["samtools", "faidx", str(fasta)],
         stdout=subprocess.DEVNULL,
@@ -225,7 +229,7 @@ def index_fasta(fasta: str | Path, write_chrom_sizes: bool = True) -> Path:
         Path: Path to the index file or chromosome sizes file.
     """
     fasta = Path(fasta)
-    pysam.faidx(str(fasta))  # creates <fasta>.fai
+    _require_pysam().faidx(str(fasta))  # creates <fasta>.fai
 
     fai = fasta.with_suffix(fasta.suffix + ".fai")
     if write_chrom_sizes:
@@ -377,8 +381,8 @@ def subsample_fasta_from_bed(
 
     fasta_handle = None
     if backend == "python":
-        assert pysam is not None
-        fasta_handle = pysam.FastaFile(str(input_FASTA))
+        pysam_mod = _require_pysam()
+        fasta_handle = pysam_mod.FastaFile(str(input_FASTA))
 
     # Open BED + output FASTA
     with input_bed.open("r") as bed, output_FASTA.open("w") as out_fasta:

src/smftools/informatics/h5ad_functions.py

@@ -9,12 +9,15 @@
 import numpy as np
 import pandas as pd
 import scipy.sparse as sp
-from pod5 import Reader
 
 from smftools.logging_utils import get_logger
+from smftools.optional_imports import require
 
 logger = get_logger(__name__)
 
+p5 = require("pod5", extra="ont", purpose="POD5 metadata")
+Reader = p5.Reader
+
 
 def add_demux_type_annotation(
     adata,