Add typing block for inspection

ultraplot/__init__.py

@@ -7,10 +7,151 @@
 
 import sys
 from pathlib import Path
-from typing import Optional
+from typing import TYPE_CHECKING, Optional
 
 from ._lazy import LazyLoader, install_module_proxy
 
+if TYPE_CHECKING:
+    # These imports are never executed at runtime, so they have zero effect on
+    # import performance. They exist solely so that type checkers (pyright, mypy)
+    # can resolve names that are otherwise provided by the lazy loader at runtime.
+    #
+    # Keep this block in sync with _LAZY_LOADING_EXCEPTIONS and every submodule's
+    # __all__ — that is the only maintenance burden.
+    import matplotlib.pyplot as pyplot
+
+    from .axes import Axes as Axes
+    from .axes import CartesianAxes as CartesianAxes
+    from .axes import ExternalAxesContainer as ExternalAxesContainer
+    from .axes import GeoAxes as GeoAxes
+    from .axes import PlotAxes as PlotAxes
+    from .axes import PolarAxes as PolarAxes
+    from .axes import ThreeAxes as ThreeAxes
+    from .colors import ColormapDatabase as ColormapDatabase
+    from .colors import ColorDatabase as ColorDatabase
+    from .colors import ContinuousColormap as ContinuousColormap
+    from .colors import DiscreteColormap as DiscreteColormap
+    from .colors import DiscreteNorm as DiscreteNorm
+    from .colors import DivergingNorm as DivergingNorm
+    from .colors import LinearSegmentedColormap as LinearSegmentedColormap
+    from .colors import LinearSegmentedNorm as LinearSegmentedNorm
+    from .colors import ListedColormap as ListedColormap
+    from .colors import PerceptualColormap as PerceptualColormap
+    from .colors import PerceptuallyUniformColormap as PerceptuallyUniformColormap
+    from .colors import SegmentedNorm as SegmentedNorm
+    from .colors import _cmap_database as colormaps
+    from .config import config_inline_backend as config_inline_backend
+    from .config import Configurator as Configurator
+    from .config import rc as rc
+    from .config import rc_matplotlib as rc_matplotlib
+    from .config import rc_ultraplot as rc_ultraplot
+    from .config import register_cmaps as register_cmaps
+    from .config import register_colors as register_colors
+    from .config import register_cycles as register_cycles
+    from .config import register_fonts as register_fonts
+    from .config import use_style as use_style
+    from .constructor import Colormap as Colormap
+    from .constructor import Cycle as Cycle
+    from .constructor import Formatter as Formatter
+    from .constructor import FORMATTERS as FORMATTERS
+    from .constructor import Locator as Locator
+    from .constructor import LOCATORS as LOCATORS
+    from .constructor import Norm as Norm
+    from .constructor import NORMS as NORMS
+    from .constructor import Proj as Proj
+    from .constructor import PROJS as PROJS
+    from .constructor import Scale as Scale
+    from .constructor import SCALES as SCALES
+    from .demos import show_channels as show_channels
+    from .demos import show_cmaps as show_cmaps
+    from .demos import show_colorspaces as show_colorspaces
+    from .demos import show_colors as show_colors
+    from .demos import show_cycles as show_cycles
+    from .demos import show_fonts as show_fonts
+    from .figure import Figure as Figure
+    from .gridspec import GridSpec as GridSpec
+    from .gridspec import SubplotGrid as SubplotGrid
+    from .legend import GeometryEntry as GeometryEntry
+    from .legend import Legend as Legend
+    from .legend import LegendEntry as LegendEntry
+    from .proj import Aitoff as Aitoff
+    from .proj import Hammer as Hammer
+    from .proj import KavrayskiyVII as KavrayskiyVII
+    from .proj import NorthPolarAzimuthalEquidistant as NorthPolarAzimuthalEquidistant
+    from .proj import NorthPolarGnomonic as NorthPolarGnomonic
+    from .proj import (
+        NorthPolarLambertAzimuthalEqualArea as NorthPolarLambertAzimuthalEqualArea,
+    )
+    from .proj import SouthPolarAzimuthalEquidistant as SouthPolarAzimuthalEquidistant
+    from .proj import SouthPolarGnomonic as SouthPolarGnomonic
+    from .proj import (
+        SouthPolarLambertAzimuthalEqualArea as SouthPolarLambertAzimuthalEqualArea,
+    )
+    from .proj import WinkelTripel as WinkelTripel
+    from .scale import CutoffScale as CutoffScale
+    from .scale import ExpScale as ExpScale
+    from .scale import FuncScale as FuncScale
+    from .scale import InverseScale as InverseScale
+    from .scale import LinearScale as LinearScale
+    from .scale import LogitScale as LogitScale
+    from .scale import LogScale as LogScale
+    from .scale import MercatorLatitudeScale as MercatorLatitudeScale
+    from .scale import PowerScale as PowerScale
+    from .scale import SineLatitudeScale as SineLatitudeScale
+    from .scale import SymmetricalLogScale as SymmetricalLogScale
+    from .text import CurvedText as CurvedText
+    from .ultralayout import ColorbarLayoutSolver as ColorbarLayoutSolver
+    from .ultralayout import compute_ultra_positions as compute_ultra_positions
+    from .ultralayout import get_grid_positions_ultra as get_grid_positions_ultra
+    from .ultralayout import is_orthogonal_layout as is_orthogonal_layout
+    from .ultralayout import UltraLayoutSolver as UltraLayoutSolver
+    from .ticker import AutoCFDatetimeFormatter as AutoCFDatetimeFormatter
+    from .ticker import AutoCFDatetimeLocator as AutoCFDatetimeLocator
+    from .ticker import AutoFormatter as AutoFormatter
+    from .ticker import CFDatetimeFormatter as CFDatetimeFormatter
+    from .ticker import DegreeFormatter as DegreeFormatter
+    from .ticker import DegreeLocator as DegreeLocator
+    from .ticker import DiscreteLocator as DiscreteLocator
+    from .ticker import FracFormatter as FracFormatter
+    from .ticker import IndexFormatter as IndexFormatter
+    from .ticker import IndexLocator as IndexLocator
+    from .ticker import LatitudeFormatter as LatitudeFormatter
+    from .ticker import LatitudeLocator as LatitudeLocator
+    from .ticker import LongitudeFormatter as LongitudeFormatter
+    from .ticker import LongitudeLocator as LongitudeLocator
+    from .ticker import SciFormatter as SciFormatter
+    from .ticker import SigFigFormatter as SigFigFormatter
+    from .ticker import SimpleFormatter as SimpleFormatter
+    from .ui import close as close
+    from .ui import figure as figure
+    from .ui import ioff as ioff
+    from .ui import ion as ion
+    from .ui import isinteractive as isinteractive
+    from .ui import show as show
+    from .ui import subplot as subplot
+    from .ui import subplots as subplots
+    from .ui import switch_backend as switch_backend
+    from .utils import arange as arange
+    from .utils import check_for_update as check_for_update
+    from .utils import edges as edges
+    from .utils import edges2d as edges2d
+    from .utils import get_colors as get_colors
+    from .utils import saturate as saturate
+    from .utils import scale_luminance as scale_luminance
+    from .utils import scale_saturation as scale_saturation
+    from .utils import set_alpha as set_alpha
+    from .utils import set_hue as set_hue
+    from .utils import set_luminance as set_luminance
+    from .utils import set_saturation as set_saturation
+    from .utils import shade as shade
+    from .utils import shift_hue as shift_hue
+    from .utils import to_hex as to_hex
+    from .utils import to_rgb as to_rgb
+    from .utils import to_rgba as to_rgba
+    from .utils import to_xyz as to_xyz
+    from .utils import to_xyza as to_xyza
+    from .utils import units as units
+
 name = "ultraplot"
 
 try:

ultraplot/tests/test_imports.py

@@ -1,8 +1,10 @@
+import ast
 import importlib.util
 import json
 import os
 import subprocess
 import sys
+from pathlib import Path
 
 import pytest
 
@@ -139,6 +141,121 @@ def test_internals_lazy_attrs():
     assert "axes.grid" in rc_matplotlib
 
 
+def _parse_type_checking_names():
+    """
+    Parse ultraplot/__init__.py and return every name declared inside the
+    ``if TYPE_CHECKING:`` block, mapped to its import origin.
+
+    Returns a dict of ``{public_name: (module, original_name)}``.
+    """
+    init_path = Path(__file__).parent.parent / "__init__.py"
+    tree = ast.parse(init_path.read_text(encoding="utf-8"))
+
+    names = {}
+    for node in tree.body:
+        if not isinstance(node, ast.If):
+            continue
+        # Match `if TYPE_CHECKING:` (handles both bare name and attribute access)
+        test = node.test
+        is_type_checking = (
+            isinstance(test, ast.Name) and test.id == "TYPE_CHECKING"
+        ) or (isinstance(test, ast.Attribute) and test.attr == "TYPE_CHECKING")
+        if not is_type_checking:
+            continue
+        for stmt in ast.walk(node):
+            if isinstance(stmt, ast.ImportFrom):
+                for alias in stmt.names:
+                    public_name = alias.asname or alias.name
+                    names[public_name] = (stmt.module, alias.name)
+            elif isinstance(stmt, ast.Import):
+                for alias in stmt.names:
+                    public_name = alias.asname or alias.name.split(".")[0]
+                    names[public_name] = (alias.name, None)
+    return names
+
+
+def test_type_checking_names_accessible_at_runtime():
+    """
+    Every name declared inside ``if TYPE_CHECKING:`` in ``__init__.py`` must
+    also be accessible at runtime via the lazy loader.
+
+    This ensures the TYPE_CHECKING block never silently drifts out of sync with
+    what the package actually exposes.
+    """
+    import ultraplot
+
+    declared = _parse_type_checking_names()
+    assert declared, "TYPE_CHECKING block is empty or could not be parsed"
+
+    missing = [name for name in declared if not hasattr(ultraplot, name)]
+    assert not missing, (
+        "Names declared in the TYPE_CHECKING block are not accessible at runtime.\n"
+        "Either remove them from the TYPE_CHECKING block or expose them via the "
+        "lazy loader / _LAZY_LOADING_EXCEPTIONS:\n"
+        + "\n".join(f"  ultraplot.{n}" for n in missing)
+    )
+
+
+def test_type_checking_block_covers_public_all():
+    """
+    Every name in ``ultraplot.__all__`` that is a public class or callable
+    should be declared in the TYPE_CHECKING block so type checkers can resolve it.
+
+    The following categories are intentionally excluded:
+
+    - Registry-derived names (e.g. ``LogNorm``) — populated dynamically from
+      matplotlib at runtime; cannot be enumerated statically.
+    - Submodule names (e.g. ``internals``) — are modules, not types.
+    - Module-level scalars defined directly in ``__init__.py`` (``name``,
+      ``version``, ``__version__``, ``setup``) — already visible to type checkers
+      without an import.
+    - Deprecated ``_rename_objs`` wrappers (e.g. ``RcConfigurator``,
+      ``inline_backend_fmt``, ``Colors``) — their runtime type is a dynamically
+      generated class/function from ``internals.warnings``; they cannot be
+      imported cleanly and are not worth exposing to type checkers.
+    """
+    import types
+    import ultraplot
+
+    declared = set(_parse_type_checking_names())
+    all_names = set(ultraplot.__all__)
+
+    # Registry-derived names — dynamically populated from matplotlib.
+    ultraplot._build_registry_map()
+    registry_names = set(ultraplot._REGISTRY_ATTRS or {})
+
+    # Submodule names — are modules, not types.
+    submodule_names = {
+        n
+        for n in all_names
+        if isinstance(getattr(ultraplot, n, None), types.ModuleType)
+    }
+
+    # Names already defined at module-level in __init__.py itself — type
+    # checkers see them directly without needing an import statement.
+    init_level = {"__version__", "version", "name", "setup"}
+
+    # Deprecated _rename_objs wrappers — their __module__ is internals.warnings
+    # because that is where the wrapper factory lives.
+    deprecated_wrappers = {
+        n
+        for n in all_names
+        if getattr(getattr(ultraplot, n, None), "__module__", "").endswith(
+            "internals.warnings"
+        )
+    }
+
+    expected = (
+        all_names - registry_names - submodule_names - init_level - deprecated_wrappers
+    )
+    missing_from_type_checking = sorted(expected - declared)
+    assert not missing_from_type_checking, (
+        "Names in ultraplot.__all__ are missing from the TYPE_CHECKING block.\n"
+        "Add them to the ``if TYPE_CHECKING:`` block in ultraplot/__init__.py:\n"
+        + "\n".join(f"  {n}" for n in missing_from_type_checking)
+    )
+
+
 def test_docstring_missing_triggers_lazy_import():
     from ultraplot.internals import docstring