[ENH] [REFACTOR] imports to improve performance of the cli. (nipoppy#281

)

* [ENH] [REFACTOR] imports to improve performance of the cli.

* [FIX] type annotation.

* [ENH] lazy import to improve cli performance.

* [FIX] `sphinxarg.ext` release `0.5.0` breaks doc build.

Issue described at sphinx-doc/sphinx-argparse#56

* [UPDATE] dependency requirement for sphinx-argparse.

.gitignore

@@ -37,3 +37,6 @@ env/
 # docs
 nipoppy_cli/docs/build
 nipoppy_cli/docs/source/schemas/*.json
+
+# OSX
+.DS_STORE

nipoppy_cli/docs/source/conf.py

@@ -95,7 +95,7 @@
     ("py:class", "argparse._SubParsersAction"),
     ("py:class", "argparse._ActionsContainer"),
     ("py:class", "StrOrPathLike"),
-    ("py:class", "nipoppy.utils.StrOrPathLike"),
+    ("py:class", "nipoppy.env.StrOrPathLike"),
     ("py:class", "typing_extensions.Self"),
 ]
 

nipoppy_cli/nipoppy/cli/parser.py

@@ -4,8 +4,7 @@
 from argparse import ArgumentParser, HelpFormatter, _ActionsContainer, _SubParsersAction
 from pathlib import Path
 
-from nipoppy.layout import DEFAULT_LAYOUT_INFO
-from nipoppy.utils import BIDS_SESSION_PREFIX, BIDS_SUBJECT_PREFIX
+from nipoppy.env import BIDS_SUBJECT_PREFIX, BIDS_SESSION_PREFIX
 
 PROGRAM_NAME = "nipoppy"
 COMMAND_INIT = "init"
@@ -202,6 +201,8 @@ def add_subparser_dicom_reorg(
     formatter_class: type[HelpFormatter] = HelpFormatter,
 ) -> ArgumentParser:
     """Add subparser for reorg command."""
+    from nipoppy.layout import DEFAULT_LAYOUT_INFO
+
     description = (
         "(Re)organize raw (DICOM) files, from the raw DICOM directory "
         f"({DEFAULT_LAYOUT_INFO.dpath_raw_imaging}) to the organized "

nipoppy_cli/nipoppy/cli/run.py

@@ -16,12 +16,6 @@
     get_global_parser,
 )
 from nipoppy.logger import add_logfile, capture_warnings, get_logger
-from nipoppy.workflows.bids_conversion import BidsConversionRunner
-from nipoppy.workflows.dataset_init import InitWorkflow
-from nipoppy.workflows.dicom_reorg import DicomReorgWorkflow
-from nipoppy.workflows.doughnut import DoughnutWorkflow
-from nipoppy.workflows.runner import PipelineRunner
-from nipoppy.workflows.tracker import PipelineTracker
 
 
 def cli(argv: Sequence[str] = None) -> None:
@@ -44,24 +38,36 @@ def cli(argv: Sequence[str] = None) -> None:
         dpath_root = args.dataset_root
 
         if command == COMMAND_INIT:
+            # Lazy import to improve performance of cli.
+            from nipoppy.workflows.dataset_init import InitWorkflow
+
             workflow = InitWorkflow(
                 dpath_root=dpath_root,
                 **workflow_kwargs,
             )
         elif command == COMMAND_DOUGHNUT:
+            # Lazy import to improve performance of cli.
+            from nipoppy.workflows.doughnut import DoughnutWorkflow
+
             workflow = DoughnutWorkflow(
                 dpath_root=dpath_root,
                 empty=args.empty,
                 regenerate=args.regenerate,
                 **workflow_kwargs,
             )
         elif command == COMMAND_DICOM_REORG:
+            # Lazy import to improve performance of cli.
+            from nipoppy.workflows.dicom_reorg import DicomReorgWorkflow
+
             workflow = DicomReorgWorkflow(
                 dpath_root=dpath_root,
                 copy_files=args.copy_files,
                 **workflow_kwargs,
             )
         elif command == COMMAND_BIDS_CONVERSION:
+            # Lazy import to improve performance of cli.
+            from nipoppy.workflows.bids_conversion import BidsConversionRunner
+
             workflow = BidsConversionRunner(
                 dpath_root=dpath_root,
                 pipeline_name=args.pipeline,
@@ -73,6 +79,9 @@ def cli(argv: Sequence[str] = None) -> None:
                 **workflow_kwargs,
             )
         elif command == COMMAND_PIPELINE_RUN:
+            # Lazy import to improve performance of cli.
+            from nipoppy.workflows.runner import PipelineRunner
+
             workflow = PipelineRunner(
                 dpath_root=dpath_root,
                 pipeline_name=args.pipeline,
@@ -84,6 +93,9 @@ def cli(argv: Sequence[str] = None) -> None:
                 **workflow_kwargs,
             )
         elif command == COMMAND_PIPELINE_TRACK:
+            # Lazy import to improve performance of cli.
+            from nipoppy.workflows.tracker import PipelineTracker
+
             workflow = PipelineTracker(
                 dpath_root=dpath_root,
                 pipeline_name=args.pipeline,

nipoppy_cli/nipoppy/config/container.py

@@ -11,7 +11,7 @@
 from pydantic import BaseModel, ConfigDict, Field
 
 from nipoppy.logger import get_logger
-from nipoppy.utils import StrOrPathLike
+from nipoppy.env import StrOrPathLike
 
 # Apptainer
 APPTAINER_BIND_FLAG = "--bind"
@@ -141,10 +141,10 @@ def add_bind_path_to_args(
     ----------
     args : list[str]
         Existing arguments
-    path_local : nipoppy.utils.StrOrPathLike
+    path_local : nipoppy.env.StrOrPathLike
         Path on disk. If this is a relative path or contains symlinks,
         it will be resolved
-    path_inside_container : Optional[nipoppy.utils.StrOrPathLike], optional
+    path_inside_container : Optional[nipoppy.env.StrOrPathLike], optional
         Path inside the container (if None, will be the same as the local path),
         by default None
     mode : str, optional

nipoppy_cli/nipoppy/config/main.py

@@ -18,11 +18,10 @@
 from nipoppy.layout import DEFAULT_LAYOUT_INFO
 from nipoppy.tabular.dicom_dir_map import DicomDirMap
 from nipoppy.utils import (
-    BIDS_SESSION_PREFIX,
-    StrOrPathLike,
     apply_substitutions_to_json,
     load_json,
 )
+from nipoppy.env import BIDS_SESSION_PREFIX, StrOrPathLike
 
 
 class Config(SchemaWithContainerConfig):
@@ -225,7 +224,7 @@ def save(self, fpath: StrOrPathLike, **kwargs):
 
         Parameters
         ----------
-        fpath : nipoppy.utils.StrOrPathLike
+        fpath : nipoppy.env.StrOrPathLike
             Path to the JSON file to write
         """
         fpath: Path = Path(fpath)

nipoppy_cli/nipoppy/env.py

@@ -0,0 +1,9 @@
+"""Variable Definitions"""
+import os
+from typing import TypeVar
+
+StrOrPathLike = TypeVar("StrOrPathLike", str, os.PathLike)
+
+# BIDS
+BIDS_SUBJECT_PREFIX = "sub-"
+BIDS_SESSION_PREFIX = "ses-"

nipoppy_cli/nipoppy/layout.py

@@ -9,10 +9,10 @@
 from nipoppy.base import Base
 from nipoppy.utils import (
     FPATH_DEFAULT_LAYOUT,
-    StrOrPathLike,
     get_pipeline_tag,
     load_json,
 )
+from nipoppy.env import StrOrPathLike
 
 
 class PathInfo(BaseModel):
@@ -144,9 +144,9 @@ def __init__(
 
         Parameters
         ----------
-        dpath_root : nipoppy.utils.StrOrPathLike
+        dpath_root : nipoppy.env.StrOrPathLike
             Path to the root directory of the dataset.
-        fpath_config : Optional[nipoppy.utils.StrOrPathLike], optional
+        fpath_config : Optional[nipoppy.env.StrOrPathLike], optional
             Path to the layout config to use, by default None.
             If None, the default layout will be used.
 

nipoppy_cli/nipoppy/logger.py

@@ -8,7 +8,7 @@
 from rich.console import Console
 from rich.logging import RichHandler
 
-from nipoppy.utils import StrOrPathLike
+from nipoppy.env import StrOrPathLike
 
 DATE_FORMAT = "[%Y-%m-%d %X]"
 FORMAT_RICH = "%(message)s"

nipoppy_cli/nipoppy/tabular/base.py

@@ -11,7 +11,8 @@
 from pydantic import BaseModel, ValidationError, model_validator
 from typing_extensions import Self
 
-from nipoppy.utils import StrOrPathLike, save_df_with_backup
+from nipoppy.utils import save_df_with_backup
+from nipoppy.env import StrOrPathLike
 
 
 class BaseTabularModel(BaseModel):

nipoppy_cli/nipoppy/tabular/doughnut.py

@@ -13,10 +13,10 @@
 from nipoppy.tabular.dicom_dir_map import DicomDirMap
 from nipoppy.tabular.manifest import Manifest, ManifestModel
 from nipoppy.utils import (
-    StrOrPathLike,
     participant_id_to_bids_participant,
     session_id_to_bids_session,
 )
+from nipoppy.env import StrOrPathLike
 
 
 class DoughnutModel(ManifestModel):

nipoppy_cli/nipoppy/utils.py

@@ -8,16 +8,12 @@
 import re
 import warnings
 from pathlib import Path
-from typing import List, Optional, Sequence, TypeVar
+from typing import List, Optional, Sequence
 
 import bids
 import pandas as pd
 
-StrOrPathLike = TypeVar("StrOrPathLike", str, os.PathLike)
-
-# BIDS
-BIDS_SUBJECT_PREFIX = "sub-"
-BIDS_SESSION_PREFIX = "ses-"
+from nipoppy.env import BIDS_SESSION_PREFIX, BIDS_SUBJECT_PREFIX, StrOrPathLike
 
 # user configs (pipeline configs, invocations, descriptors)
 TEMPLATE_REPLACE_PATTERN = re.compile("\\[\\[NIPOPPY\\_(.*?)\\]\\]")
@@ -196,7 +192,7 @@ def load_json(fpath: StrOrPathLike, **kwargs) -> dict:
 
     Parameters
     ----------
-    fpath : nipoppy.utils.StrOrPathLike
+    fpath : nipoppy.env.StrOrPathLike
         Path to the JSON file
     **kwargs :
         Keyword arguments to pass to json.load
@@ -222,7 +218,7 @@ def save_json(obj: dict, fpath: StrOrPathLike, **kwargs):
     ----------
     obj : dict
         The JSON object
-    fpath : nipoppy.utils.StrOrPathLike
+    fpath : nipoppy.env.StrOrPathLike
         Path to the JSON file to write
     indent : int, optional
         Indentation level, by default 4
@@ -265,7 +261,7 @@ def save_df_with_backup(
     ----------
     df : pd.DataFrame
         The dataframe to save
-    fpath_symlink : nipoppy.utils.StrOrPathLike
+    fpath_symlink : nipoppy.env.StrOrPathLike
         The path to the symlink
     dname_backups : Optional[str], optional
         The directory where the timestamped backup file should be written

nipoppy_cli/nipoppy/workflows/base.py

@@ -21,7 +21,8 @@
 from nipoppy.tabular.dicom_dir_map import DicomDirMap
 from nipoppy.tabular.doughnut import Doughnut, generate_doughnut
 from nipoppy.tabular.manifest import Manifest
-from nipoppy.utils import StrOrPathLike, add_path_timestamp, process_template_str
+from nipoppy.utils import add_path_timestamp, process_template_str
+from nipoppy.env import StrOrPathLike
 
 LOG_SUFFIX = ".log"
 
@@ -51,7 +52,7 @@ def __init__(
 
         Parameters
         ----------
-        dpath_root : nipoppy.utils.StrOrPathLike
+        dpath_root : nipoppy.env.StrOrPathLike
             Path the the root directory of the dataset.
         name : str
             Name of the workflow, used for logging.

nipoppy_cli/nipoppy/workflows/bids_conversion.py

@@ -8,7 +8,7 @@
 from typing import Optional
 
 from nipoppy.config.pipeline import BidsPipelineConfig
-from nipoppy.utils import StrOrPathLike
+from nipoppy.env import StrOrPathLike
 from nipoppy.workflows.runner import PipelineRunner
 
 

nipoppy_cli/nipoppy/workflows/dataset_init.py

@@ -10,8 +10,8 @@
     DPATH_TRACKER_CONFIGS,
     FPATH_SAMPLE_CONFIG,
     FPATH_SAMPLE_MANIFEST,
-    StrOrPathLike,
 )
+from nipoppy.env import StrOrPathLike
 from nipoppy.workflows.base import BaseWorkflow
 
 

nipoppy_cli/nipoppy/workflows/dicom_reorg.py

@@ -9,10 +9,10 @@
 
 from nipoppy.tabular.doughnut import update_doughnut
 from nipoppy.utils import (
-    StrOrPathLike,
     participant_id_to_bids_participant,
     session_id_to_bids_session,
 )
+from nipoppy.env import StrOrPathLike
 from nipoppy.workflows.base import BaseWorkflow
 
 

nipoppy_cli/nipoppy/workflows/doughnut.py

@@ -5,7 +5,7 @@
 from typing import Optional
 
 from nipoppy.tabular.doughnut import Doughnut, generate_doughnut, update_doughnut
-from nipoppy.utils import StrOrPathLike
+from nipoppy.env import StrOrPathLike
 from nipoppy.workflows.base import BaseWorkflow
 
 

nipoppy_cli/nipoppy/workflows/pipeline.py

@@ -11,6 +11,7 @@
 from typing import Optional
 
 import bids
+from nipoppy.env import BIDS_SESSION_PREFIX, BIDS_SUBJECT_PREFIX
 from pydantic import ValidationError
 
 from nipoppy.config.boutiques import (
@@ -19,9 +20,6 @@
 )
 from nipoppy.config.pipeline import ProcPipelineConfig
 from nipoppy.utils import (
-    BIDS_SESSION_PREFIX,
-    BIDS_SUBJECT_PREFIX,
-    StrOrPathLike,
     add_pybids_ignore_patterns,
     check_participant_id,
     check_session_id,
@@ -32,6 +30,7 @@
     process_template_str,
     session_id_to_bids_session,
 )
+from nipoppy.env import StrOrPathLike
 from nipoppy.workflows.base import BaseWorkflow
 
 

nipoppy_cli/nipoppy/workflows/runner.py

@@ -10,7 +10,7 @@
 from nipoppy.config.boutiques import BoutiquesConfig
 from nipoppy.config.container import ContainerConfig, prepare_container
 from nipoppy.tabular.bagel import Bagel
-from nipoppy.utils import StrOrPathLike
+from nipoppy.env import StrOrPathLike
 from nipoppy.workflows.pipeline import BasePipelineWorkflow
 
 

nipoppy_cli/nipoppy/workflows/tracker.py

@@ -7,7 +7,8 @@
 
 from nipoppy.config.tracker import TrackerConfig, check_tracker_configs
 from nipoppy.tabular.bagel import Bagel
-from nipoppy.utils import StrOrPathLike, load_json
+from nipoppy.utils import load_json
+from nipoppy.env import StrOrPathLike
 from nipoppy.workflows.pipeline import BasePipelineWorkflow
 
 

nipoppy_cli/pyproject.toml

@@ -42,7 +42,7 @@ doc = [
     "furo>=2024.1.29",
     "pygments-csv-lexer>=0.1.3",
     "sphinx>=7.2.6",
-    "sphinx-argparse>=0.4.0",
+    "sphinx-argparse>=0.4.0, !=0.5.0",
     "sphinx-autoapi==3.0.0",
     "sphinx-copybutton>=0.5.2",
     "sphinx-jsonschema>=1.19.1",

nipoppy_cli/tests/conftest.py

@@ -16,10 +16,10 @@
 from nipoppy.tabular.doughnut import Doughnut
 from nipoppy.tabular.manifest import Manifest
 from nipoppy.utils import (
-    StrOrPathLike,
     participant_id_to_bids_participant,
     session_id_to_bids_session,
 )
+from nipoppy.env import StrOrPathLike
 
 FPATH_CONFIG = "global_config.json"
 FPATH_MANIFEST = "manifest.csv"

nipoppy_cli/tests/test_tabular_doughnut.py

@@ -7,7 +7,7 @@
 
 from nipoppy.tabular.dicom_dir_map import DicomDirMap
 from nipoppy.tabular.doughnut import Doughnut, generate_doughnut, update_doughnut
-from nipoppy.utils import StrOrPathLike
+from nipoppy.env import StrOrPathLike
 
 from .conftest import DPATH_TEST_DATA, check_doughnut, prepare_dataset