refactor(core): centralize optional-dependency

Author: mkumar73

deeptab/core/observability.py

@@ -300,12 +300,9 @@ def build_structlog_logger(config: ObservabilityConfig, run_dir: str | None = No
     ImportError
         If ``structlog`` is not installed, with an actionable install hint.
     """
-    try:
-        import structlog  # type: ignore[import-untyped]
-    except ImportError as exc:
-        raise ImportError(
-            "structlog is required when structured_logging=True. Install it with: pip install 'deeptab[logs]'"
-        ) from exc
+    from deeptab.core.optional_deps import require_structlog
+
+    structlog = require_structlog()
 
     import json
     import os
@@ -458,16 +455,15 @@ def build_lightning_loggers(
     """
     import os
 
+    from deeptab.core.optional_deps import require_mlflow, require_tensorboard
+
     loggers: list[Any] = []
 
     for tracker in config.experiment_trackers:
         if tracker == "mlflow":
-            try:
-                from lightning.pytorch.loggers import MLFlowLogger
-            except ImportError as exc:
-                raise ImportError(
-                    "MLflow logging requires the mlflow package. Install it with: pip install 'deeptab[mlflow]'"
-                ) from exc
+            require_mlflow()
+            from lightning.pytorch.loggers import MLFlowLogger
+
             # Ensure the artifact location directory exists
             if config.mlflow_artifact_location:
                 os.makedirs(config.mlflow_artifact_location, exist_ok=True)
@@ -482,13 +478,9 @@ def build_lightning_loggers(
             )
 
         elif tracker == "tensorboard":
-            try:
-                from lightning.pytorch.loggers import TensorBoardLogger
-            except ImportError as exc:
-                raise ImportError(
-                    "TensorBoard logging requires the tensorboard package. "
-                    "Install it with: pip install 'deeptab[tensorboard]'"
-                ) from exc
+            require_tensorboard()
+            from lightning.pytorch.loggers import TensorBoardLogger
+
             loggers.append(
                 TensorBoardLogger(
                     save_dir=config.tensorboard_save_dir,

deeptab/core/optional_deps.py

@@ -0,0 +1,86 @@
+"""Guards for optional third-party dependencies.
+
+DeepTab keeps its default install lightweight: structured logging and the
+experiment-tracking backends are shipped as optional extras. Each helper here
+imports one optional dependency on demand and raises a clear, actionable
+:class:`ImportError` — pointing at the matching ``pip install 'deeptab[...]'``
+command — when the package is not installed.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from types import ModuleType
+
+
+def require_structlog() -> ModuleType:
+    """Return the :mod:`structlog` module.
+
+    Returns
+    -------
+    module
+        The imported ``structlog`` module.
+
+    Raises
+    ------
+    ImportError
+        If ``structlog`` is not installed, with an actionable install hint.
+    """
+    try:
+        import structlog  # type: ignore[import-untyped]
+    except ImportError as exc:
+        raise ImportError(
+            "Structured logging requires the optional 'structlog' dependency. "
+            "Install it with: pip install 'deeptab[logs]'"
+        ) from exc
+
+    return structlog
+
+
+def require_mlflow() -> ModuleType:
+    """Return the :mod:`mlflow` module.
+
+    Returns
+    -------
+    module
+        The imported ``mlflow`` module.
+
+    Raises
+    ------
+    ImportError
+        If ``mlflow`` is not installed, with an actionable install hint.
+    """
+    try:
+        import mlflow  # type: ignore[import-untyped]
+    except ImportError as exc:
+        raise ImportError(
+            "MLflow tracking requires the optional 'mlflow' dependency. Install it with: pip install 'deeptab[mlflow]'"
+        ) from exc
+
+    return mlflow
+
+
+def require_tensorboard() -> Any:
+    """Return ``torch.utils.tensorboard.SummaryWriter``.
+
+    Returns
+    -------
+    type
+        The ``SummaryWriter`` class from ``torch.utils.tensorboard``.
+
+    Raises
+    ------
+    ImportError
+        If ``tensorboard`` is not installed, with an actionable install hint.
+    """
+    try:
+        from torch.utils.tensorboard import SummaryWriter
+    except ImportError as exc:
+        raise ImportError(
+            "TensorBoard logging requires the optional 'tensorboard' dependency. "
+            "Install it with: pip install 'deeptab[tensorboard]'"
+        ) from exc
+
+    return SummaryWriter

tests/test_observability.py

@@ -158,32 +158,17 @@ def test_build_lightning_loggers_unknown_tracker_raises():
 
 def test_build_lightning_loggers_mlflow_absent(monkeypatch):
     """ImportError with install hint when mlflow is not installed."""
-    # Simulate mlflow being absent by blocking its import inside Lightning
-    real_import = __builtins__.__import__ if hasattr(__builtins__, "__import__") else __import__  # type: ignore[attr-defined]
-
-    def _block_mlflow(name, *args, **kwargs):
-        if "MLFlowLogger" in name or (len(args) >= 3 and "MLFlowLogger" in str(args[2])):
-            raise ImportError("No module named 'mlflow'")
-        return real_import(name, *args, **kwargs)
-
-    # Use monkeypatch on the lightning loggers module directly
-    mock_module = MagicMock()
-    mock_module.MLFlowLogger.side_effect = ImportError("No module named 'mlflow'")
-
-    import lightning.pytorch.loggers as lpl
-
-    original_MLFlowLogger = getattr(lpl, "MLFlowLogger", None)
-
-    # Patch lightning.pytorch.loggers so that importing MLFlowLogger raises
-    monkeypatch.setitem(sys.modules, "lightning.pytorch.loggers", None)  # type: ignore[arg-type]
+    # The guard checks the actual ``mlflow`` package, so simulate its absence.
+    monkeypatch.setitem(sys.modules, "mlflow", None)  # type: ignore[arg-type]
     cfg = ObservabilityConfig(experiment_trackers=["mlflow"])
     with pytest.raises(ImportError, match="pip install 'deeptab\\[mlflow\\]'"):
         build_lightning_loggers(cfg)
 
 
 def test_build_lightning_loggers_tensorboard_absent(monkeypatch):
     """ImportError with install hint when tensorboard is not installed."""
-    monkeypatch.setitem(sys.modules, "lightning.pytorch.loggers", None)  # type: ignore[arg-type]
+    # The guard checks ``torch.utils.tensorboard``, so simulate its absence.
+    monkeypatch.setitem(sys.modules, "torch.utils.tensorboard", None)  # type: ignore[arg-type]
     cfg = ObservabilityConfig(experiment_trackers=["tensorboard"])
     with pytest.raises(ImportError, match="pip install 'deeptab\\[tensorboard\\]'"):
         build_lightning_loggers(cfg)