Skip to content

add stub-driven lazy imports using lazy_loader #3459

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
deruyter92 wants to merge 1 commit into from
Draft

add stub-driven lazy imports using lazy_loader #3459

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions .pre-commit-config.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@ repos:
stages: [pre-commit, manual]
- id: name-tests-test
args: [--pytest-test-first]
exclude: ^tests/typing/
stages: [pre-commit, manual]
- id: check-json
stages: [pre-commit, manual]
Expand Down
288 changes: 49 additions & 239 deletions deeplabcut/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,272 +13,82 @@

import logging
import os
from importlib import import_module
from typing import Any
import warnings

import lazy_loader as lazy

from deeplabcut.core.deprecation import DLCDeprecationWarning

from .version import VERSION, __version__

logger = logging.getLogger(__name__)

# DEBUG="", "0", "false", "no" -> False
DEBUG = os.environ.get("DEBUG", "").strip().lower() not in {"", "0", "false", "no"}

from .version import VERSION, __version__

if DEBUG:
logger.debug("Loading DLC %s", VERSION)

# DeepLabCut deprecation warnings are shown only once per message instance.
import warnings

from deeplabcut.core.deprecation import DLCDeprecationWarning

warnings.filterwarnings("once", category=DLCDeprecationWarning)

# -----------------------------------------------------------------------------
# Always-available public API
# Stub-driven lazy loading
# -----------------------------------------------------------------------------
# ``deeplabcut/__init__.pyi`` is the single declarative source of truth for the
# top-level public API. ``lazy_loader.attach_stub`` reads it at runtime to
# install ``__getattr__``, ``__dir__``, and ``__all__``, so each implementation
# module is imported only when its top-level attribute is first accessed.
# -----------------------------------------------------------------------------

from .core.engine import Engine
from .create_project import (
add_new_videos,
create_new_project,
create_new_project_3d,
create_pretrained_human_project,
load_demo_data,
)
from .generate_training_dataset import (
adddatasetstovideolistandviceversa,
check_labels,
comparevideolistsanddatafolders,
create_multianimaltraining_dataset,
create_training_dataset,
create_training_dataset_from_existing_split,
create_training_model_comparison,
dropannotationfileentriesduetodeletedimages,
dropduplicatesinannotatinfiles,
dropimagesduetolackofannotation,
dropunlabeledframes,
extract_frames,
mergeandsplit,
)
from .pose_estimation_3d import (
calibrate_cameras,
check_undistortion,
create_labeled_video_3d,
triangulate,
)
from .utils import (
analyze_videos_converth5_to_csv,
analyze_videos_converth5_to_nwb,
auxfun_videos,
auxiliaryfunctions,
convert2_maDLC,
convertcsv2h5,
create_labeled_video,
create_video_with_all_detections,
plot_trajectories,
)
from .utils.auxfun_videos import (
CropVideo,
DownSampleVideo,
ShortenVideo,
check_video_integrity,
collect_video_paths,
)
_lazy_getattr, __dir__, __all__ = lazy.attach_stub(__name__, __file__)

# -----------------------------------------------------------------------------
# Optional / lazy public API
# Optional-dependency diagnostics
# -----------------------------------------------------------------------------
# These names are part of the public API, but importing them may require
# optional GUI or torch dependencies, so we lazy load them.
#
# Example:
# import deeplabcut as dlc
# dlc.launch_dlc() # imports GUI code lazily
# dlc.transformer_reID(...) # imports torch-dependent code lazily
# A plain ``attach_stub`` raises ``ModuleNotFoundError`` when a GUI or PyTorch
# tracking module is unavailable. Translate only those into actionable
# ``ImportError`` messages and leave unrelated import failures untouched.
# -----------------------------------------------------------------------------

_OPTIONAL_EXPORTS: dict[str, tuple[str, str]] = {
# GUI
"launch_dlc": (".gui.launch_script", "launch_dlc"),
"label_frames": (".gui.tabs.label_frames", "label_frames"),
"refine_labels": (".gui.tabs.label_frames", "refine_labels"),
"refine_tracklets": (".gui.tracklet_toolbox", "refine_tracklets"),
"SkeletonBuilder": (".gui.widgets", "SkeletonBuilder"),
# Optional torch feature
"transformer_reID": (".pose_tracking_pytorch", "transformer_reID"),
}

# API exports are lazily loaded from the pose_estimation API facade.
_API_EXPORTS_MAP: dict[str, tuple[str, str]] = {
"analyze_images": (".api.pose_estimation", "analyze_images"),
"analyze_time_lapse_frames": (".api.pose_estimation", "analyze_time_lapse_frames"),
"analyze_videos": (".api.pose_estimation", "analyze_videos"),
"convert_detections2tracklets": (".api.pose_estimation", "convert_detections2tracklets"),
"create_pretrained_project": (".api.create_project", "create_pretrained_project"),
"create_tracking_dataset": (".api.pose_estimation", "create_tracking_dataset"),
"evaluate_network": (".api.pose_estimation", "evaluate_network"),
"export_model": (".api.pose_estimation", "export_model"),
"extract_maps": (".api.pose_estimation", "extract_maps"),
"extract_save_all_maps": (".api.pose_estimation", "extract_save_all_maps"),
"return_evaluate_network_data": (".api.pose_estimation", "return_evaluate_network_data"),
"return_train_network_path": (".api.pose_estimation", "return_train_network_path"),
"train_network": (".api.pose_estimation", "train_network"),
"visualize_locrefs": (".api.pose_estimation", "visualize_locrefs"),
"visualize_paf": (".api.pose_estimation", "visualize_paf"),
"visualize_scoremaps": (".api.pose_estimation", "visualize_scoremaps"),
"analyzeskeleton": (".api.post_processing", "analyzeskeleton"),
"filterpredictions": (".api.post_processing", "filterpredictions"),
"extract_outlier_frames": (".api.refine_training", "extract_outlier_frames"),
"find_outliers_in_raw_data": (".api.refine_training", "find_outliers_in_raw_data"),
"merge_datasets": (".api.refine_training", "merge_datasets"),
"stitch_tracklets": (".api.refine_training", "stitch_tracklets"),
"video_inference_superanimal": (".api.modelzoo_inference", "video_inference_superanimal"),
}


def __getattr__(name: str) -> Any:
"""Lazily load optional public exports and API exports."""
# Check API exports first (always available, lightweight import)
if name in _API_EXPORTS_MAP:
module_name, attr_name = _API_EXPORTS_MAP[name]
module = import_module(module_name, package=__name__)
value = getattr(module, attr_name)
globals()[name] = value
return value

if name not in _OPTIONAL_EXPORTS:
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

module_name, attr_name = _OPTIONAL_EXPORTS[name]

try:
module = import_module(module_name, package=__name__)
value = getattr(module, attr_name)
except (ModuleNotFoundError, ImportError) as exc:
if name in {
"launch_dlc",
"label_frames",
"refine_labels",
"refine_tracklets",
"SkeletonBuilder",
}:
raise AttributeError(
f"{name!r} is unavailable because DeepLabCut was loaded without GUI dependencies."
) from exc

if name == "transformer_reID":
raise AttributeError(
f"{name!r} is unavailable because the PyTorch-based tracking dependencies are not installed."
) from exc

raise AttributeError(f"module {__name__!r} has no attribute {name!r}") from exc

# Cache the resolved object so future access is fast
globals()[name] = value
return value


def __dir__() -> list[str]:
"""Improve IDE / autocomplete discoverability."""
return sorted(set(globals()) | set(__all__))

_GUI_EXPORTS = frozenset(
{
"launch_dlc",
"label_frames",
"refine_labels",
"refine_tracklets",
"SkeletonBuilder",
}
)

# -----------------------------------------------------------------------------
# Public API
# -----------------------------------------------------------------------------
_TORCH_EXPORTS = frozenset({"transformer_reID"})

@C-Achard C-Achard Aug 21, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

How do you rate this in terms of maintainability?
Could we parse from the pyproject somehow?
Or make a shared helper that is called by all e.g. GUI modules that tries to import gui deps,
and raises directly an import error asking to install the optional dep? This way each module with optional deps has clear guidance messages, we have a centralized file with all of those at a glance, and it only couples it to local source code imports without risking drift with the pyproject


_VERSION_EXPORTS = [
"__version__",
"VERSION",
"DEBUG",
]
_GUI_DEPENDENCY_MODULES = frozenset({"PySide6", "napari", "qdarkstyle"})

@C-Achard C-Achard Aug 21, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should this be broader? What is napari-deeplabcut is missing?

_TORCH_DEPENDENCY_MODULES = frozenset({"torch", "torchvision"})

_CORE_EXPORTS = [
"Engine",
]

_PROJECT_EXPORTS = [
"add_new_videos",
"create_new_project",
"create_new_project_3d",
"create_pretrained_human_project",
"load_demo_data",
]
def _is_missing_gui_dependency(exc: ModuleNotFoundError) -> bool:
"""Return True if ``exc`` is caused by a missing GUI dependency."""
name = getattr(exc, "name", None)
return isinstance(name, str) and name.split(".")[0] in _GUI_DEPENDENCY_MODULES

_DATASET_EXPORTS = [
"adddatasetstovideolistandviceversa",
"check_labels",
"comparevideolistsanddatafolders",
"create_multianimaltraining_dataset",
"create_training_dataset",
"create_training_dataset_from_existing_split",
"create_training_model_comparison",
"dropannotationfileentriesduetodeletedimages",
"dropduplicatesinannotatinfiles",
"dropimagesduetolackofannotation",
"dropunlabeledframes",
"extract_frames",
"mergeandsplit",
]

_API_EXPORTS = [
"analyze_images",
"analyze_time_lapse_frames",
"analyze_videos",
"convert_detections2tracklets",
"create_tracking_dataset",
"evaluate_network",
"export_model",
"extract_maps",
"extract_save_all_maps",
"return_evaluate_network_data",
"return_train_network_path",
"train_network",
"visualize_locrefs",
"visualize_paf",
"visualize_scoremaps",
"analyzeskeleton",
"create_pretrained_project",
"filterpredictions",
"extract_outlier_frames",
"find_outliers_in_raw_data",
"merge_datasets",
"stitch_tracklets",
"video_inference_superanimal",
]
def _is_missing_torch_dependency(exc: ModuleNotFoundError) -> bool:
"""Return True if ``exc`` is caused by a missing PyTorch dependency."""
name = getattr(exc, "name", None)
return isinstance(name, str) and name.split(".")[0] in _TORCH_DEPENDENCY_MODULES

_UTIL_EXPORTS = [
"analyze_videos_converth5_to_csv",
"analyze_videos_converth5_to_nwb",
"auxfun_videos",
"auxiliaryfunctions",
"convert2_maDLC",
"convertcsv2h5",
"create_labeled_video",
"create_video_with_all_detections",
"plot_trajectories",
"CropVideo",
"DownSampleVideo",
"ShortenVideo",
"check_video_integrity",
]

_THREE_D_EXPORTS = [
"calibrate_cameras",
"check_undistortion",
"create_labeled_video_3d",
"triangulate",
]
def __getattr__(name: str):
try:
return _lazy_getattr(name)
except ModuleNotFoundError as exc:
if name in _GUI_EXPORTS and _is_missing_gui_dependency(exc):
raise ImportError(
f"{name!r} requires the DeepLabCut GUI dependencies. Install the supported GUI extra."
) from exc

_OPTIONAL_API_EXPORTS = list(_OPTIONAL_EXPORTS)
if name in _TORCH_EXPORTS and _is_missing_torch_dependency(exc):
raise ImportError(f"{name!r} requires the PyTorch tracking dependencies.") from exc

__all__ = (
_VERSION_EXPORTS
+ _CORE_EXPORTS
+ _PROJECT_EXPORTS
+ _DATASET_EXPORTS
+ _API_EXPORTS
+ _UTIL_EXPORTS
+ _THREE_D_EXPORTS
+ _OPTIONAL_API_EXPORTS
)
raise
Comment on lines +82 to +94

@C-Achard C-Achard Aug 21, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe good to mention: this changes e.g. hasattr(deeplabcut, "launch_dlc") from returning False to raising.
I would assume this is fine in most cases but still a bit of a sneaky change

Loading