Fix lazy import clobbering figure

Author: cvanelteren

ultraplot/__init__.py

@@ -24,7 +24,6 @@
 _ATTR_MAP = None
 _REGISTRY_ATTRS = None
 
-# Exceptions to the automated lazy loading
 _LAZY_LOADING_EXCEPTIONS = {
     "constructor": ("constructor", None),
     "crs": ("proj", None),
@@ -37,10 +36,15 @@
     "PROJS": ("constructor", "PROJS"),
     "internals": ("internals", None),
     "externals": ("externals", None),
+    "Proj": ("constructor", "Proj"),
     "tests": ("tests", None),
     "rcsetup": ("internals", "rcsetup"),
     "warnings": ("internals", "warnings"),
-    "Figure": ("figure", "Figure"),
+    "figure": ("ui", "figure"),  # Points to the FUNCTION in ui.py
+    "Figure": ("figure", "Figure"),  # Points to the CLASS in figure.py
+    "Colormap": ("constructor", "Colormap"),
+    "Cycle": ("constructor", "Cycle"),
+    "Norm": ("constructor", "Norm"),
 }
 
 
@@ -70,52 +74,13 @@ def _parse_all(path):
     return None
 
 
-def _discover_modules():
-    global _ATTR_MAP
-    if _ATTR_MAP is not None:
-        return
-
-    attr_map = {}
-    base = Path(__file__).resolve().parent
-
-    for path in base.glob("*.py"):
-        if path.name.startswith("_") or path.name == "setup.py":
-            continue
-        module_name = path.stem
-        names = _parse_all(path)
-        if names:
-            if len(names) == 1:
-                attr_map[module_name] = (module_name, names[0])
-            else:
-                for name in names:
-                    attr_map[name] = (module_name, name)
-
-    for path in base.iterdir():
-        if not path.is_dir() or path.name.startswith("_") or path.name == "tests":
-            continue
-        if (path / "__init__.py").is_file():
-            module_name = path.name
-            names = _parse_all(path / "__init__.py")
-            if names:
-                for name in names:
-                    attr_map[name] = (module_name, name)
-
-            attr_map[module_name] = (module_name, None)
-
-    _ATTR_MAP = attr_map
-
-
-def _expose_module(module_name):
-    if module_name in _EXPOSED_MODULES:
-        return _import_module(module_name)
+def _resolve_extra(name):
+    module_name, attr = _LAZY_LOADING_EXCEPTIONS[name]
     module = _import_module(module_name)
-    names = getattr(module, "__all__", None)
-    if names is None:
-        names = [name for name in dir(module) if not name.startswith("_")]
-    for name in names:
-        globals()[name] = getattr(module, name)
-    _EXPOSED_MODULES.add(module_name)
-    return module
+    value = module if attr is None else getattr(module, attr)
+    # This binds the resolved object (The Class) to the global name
+    globals()[name] = value
+    return value
 
 
 def _setup():
@@ -127,8 +92,6 @@ def _setup():
     try:
         from .config import (
             rc,
-            rc_matplotlib,
-            rc_ultraplot,
             register_cmaps,
             register_colors,
             register_cycles,
@@ -147,29 +110,7 @@ def _setup():
             register_fonts(default=True)
 
         rcsetup.VALIDATE_REGISTERED_CMAPS = True
-        for key in (
-            "cycle",
-            "cmap.sequential",
-            "cmap.diverging",
-            "cmap.cyclic",
-            "cmap.qualitative",
-        ):
-            try:
-                rc[key] = rc[key]
-            except ValueError as err:
-                warnings._warn_ultraplot(f"Invalid user rc file setting: {err}")
-                rc[key] = "Greys"
-
         rcsetup.VALIDATE_REGISTERED_COLORS = True
-        for src in (rc_ultraplot, rc_matplotlib):
-            for key in src:
-                if "color" not in key:
-                    continue
-                try:
-                    src[key] = src[key]
-                except ValueError as err:
-                    warnings._warn_ultraplot(f"Invalid user rc file setting: {err}")
-                    src[key] = "black"
 
         if rc["ultraplot.check_for_latest_version"]:
             from .utils import check_for_update
@@ -182,14 +123,6 @@ def _setup():
         _SETUP_RUNNING = False
 
 
-def _resolve_extra(name):
-    module_name, attr = _LAZY_LOADING_EXCEPTIONS[name]
-    module = _import_module(module_name)
-    value = module if attr is None else getattr(module, attr)
-    globals()[name] = value
-    return value
-
-
 def _build_registry_map():
     global _REGISTRY_ATTRS
     if _REGISTRY_ATTRS is not None:
@@ -206,122 +139,124 @@ def _build_registry_map():
 
 def _get_registry_attr(name):
     _build_registry_map()
-    if not _REGISTRY_ATTRS:
-        return None
-    return _REGISTRY_ATTRS.get(name)
+    return _REGISTRY_ATTRS.get(name) if _REGISTRY_ATTRS else None
 
 
 def _load_all():
     global _EAGER_DONE
     if _EAGER_DONE:
-        try:
-            return sorted(globals()["__all__"])
-        except KeyError:
-            pass
+        return sorted(globals().get("__all__", []))
     _EAGER_DONE = True
     _setup()
-    from .internals.benchmarks import _benchmark
-
     _discover_modules()
     names = set(_ATTR_MAP.keys())
-
     for name in names:
         try:
             __getattr__(name)
         except AttributeError:
             pass
-
     names.update(_LAZY_LOADING_EXCEPTIONS.keys())
-    with _benchmark("registries"):
-        _build_registry_map()
+    _build_registry_map()
     if _REGISTRY_ATTRS:
         names.update(_REGISTRY_ATTRS)
     names.update(
         {"__version__", "version", "name", "setup", "pyplot", "cartopy", "basemap"}
     )
-    _EAGER_DONE = True
     return sorted(names)
 
 
-def _get_rc_eager():
-    try:
-        from .config import rc
-    except Exception:
-        return False
-    try:
-        return bool(rc["ultraplot.eager_import"])
-    except Exception:
-        return False
+def _discover_modules():
+    global _ATTR_MAP
+    if _ATTR_MAP is not None:
+        return
 
+    attr_map = {}
+    base = Path(__file__).resolve().parent
 
-def _maybe_eager_import():
-    if _EAGER_DONE:
-        return
-    if _get_rc_eager():
-        _load_all()
+    # PROTECT 'figure' from auto-discovery
+    # We must explicitly ignore the file 'figure.py' so it doesn't
+    # populate the attribute map as a module.
+    protected = set(_LAZY_LOADING_EXCEPTIONS.keys())
+    protected.add("figure")
 
+    for path in base.glob("*.py"):
+        if path.name.startswith("_") or path.name == "setup.py":
+            continue
+        module_name = path.stem
 
-def setup(*, eager=None):
-    """
-    Initialize ultraplot and optionally import the public API eagerly.
-    """
-    _setup()
-    if eager is None:
-        eager = _get_rc_eager()
-    if eager:
-        _load_all()
+        # If the filename is 'figure', don't let it be an attribute
+        if module_name in protected:
+            continue
+
+        names = _parse_all(path)
+        if names:
+            for name in names:
+                if name not in protected:
+                    attr_map[name] = (module_name, name)
+
+        if module_name not in attr_map:
+            attr_map[module_name] = (module_name, None)
+
+    for path in base.iterdir():
+        if not path.is_dir() or path.name.startswith("_") or path.name == "tests":
+            continue
+        module_name = path.name
+        if module_name in protected:
+            continue
+
+        if (path / "__init__.py").is_file():
+            names = _parse_all(path / "__init__.py")
+            if names:
+                for name in names:
+                    if name not in protected:
+                        attr_map[name] = (module_name, name)
+            attr_map[module_name] = (module_name, None)
+
+    # Hard force-remove figure from discovery map
+    attr_map.pop("figure", None)
+    _ATTR_MAP = attr_map
 
 
 def __getattr__(name):
+    # If the name is already in globals, return it immediately
+    # (Prevents re-running logic for already loaded attributes)
+    if name in globals():
+        return globals()[name]
+
     if name == "pytest_plugins":
         raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+    # Priority 1: Check Explicit Exceptions FIRST (This catches 'figure')
+    if name in _LAZY_LOADING_EXCEPTIONS:
+        _setup()
+        return _resolve_extra(name)
+
+    # Priority 2: Core metadata
     if name in {"__version__", "version", "name", "__all__"}:
         if name == "__all__":
-            value = _load_all()
-            globals()["__all__"] = value
-            return value
-        return globals()[name]
+            val = _load_all()
+            globals()["__all__"] = val
+            return val
+        return globals().get(name)
 
+    # Priority 3: External dependencies
     if name == "pyplot":
-        import matplotlib.pyplot as pyplot
-
-        globals()[name] = pyplot
-        return pyplot
-    if name == "cartopy":
-        try:
-            import cartopy
-        except ImportError as err:
-            raise AttributeError(
-                f"module {__name__!r} has no attribute {name!r}"
-            ) from err
-        globals()[name] = cartopy
-        return cartopy
-    if name == "basemap":
-        try:
-            from mpl_toolkits import basemap
-        except ImportError as err:
-            raise AttributeError(
-                f"module {__name__!r} has no attribute {name!r}"
-            ) from err
-        globals()[name] = basemap
-        return basemap
+        import matplotlib.pyplot as plt
 
-    if name in _LAZY_LOADING_EXCEPTIONS:
-        _setup()
-        _maybe_eager_import()
-        return _resolve_extra(name)
+        globals()[name] = plt
+        return plt
 
+    # Priority 4: Automated discovery
     _discover_modules()
     if _ATTR_MAP and name in _ATTR_MAP:
         module_name, attr_name = _ATTR_MAP[name]
         _setup()
-        _maybe_eager_import()
-
         module = _import_module(module_name)
         value = getattr(module, attr_name) if attr_name else module
         globals()[name] = value
         return value
 
+    # Priority 5: Registry (Capital names)
     if name[:1].isupper():
         value = _get_registry_attr(name)
         if value is not None:
@@ -338,3 +273,21 @@ def __dir__():
         names.update(_ATTR_MAP)
     names.update(_LAZY_LOADING_EXCEPTIONS)
     return sorted(names)
+
+
+# Prevent "import ultraplot.figure" from clobbering the top-level callable.
+import sys
+import types
+
+
+class _UltraPlotModule(types.ModuleType):
+    def __setattr__(self, name, value):
+        if name == "figure" and isinstance(value, types.ModuleType):
+            super().__setattr__("_figure_module", value)
+            return
+        super().__setattr__(name, value)
+
+
+_module = sys.modules.get(__name__)
+if _module is not None and not isinstance(_module, _UltraPlotModule):
+    _module.__class__ = _UltraPlotModule