fix: cache lazy imports correctly and expose diagram module

- Cache lazy imports in globals() to override the submodule that
  importlib automatically sets on the parent module
- Add dj.diagram to lazy modules (returns module for diagram_active access)
- Add tests for cli callable and diagram module access

Co-Authored-By: Claude Opus 4.5 <[email protected]>

Author: dimitri-yatsenko

src/datajoint/__init__.py

@@ -96,6 +96,7 @@
     "Diagram": (".diagram", "Diagram"),
     "Di": (".diagram", "Diagram"),
     "ERD": (".diagram", "Diagram"),
+    "diagram": (".diagram", None),  # Return the module itself
     # kill imports pymysql via connection
     "kill": (".admin", "kill"),
     # cli imports click
@@ -110,5 +111,10 @@ def __getattr__(name: str):
         import importlib
 
         module = importlib.import_module(module_path, __package__)
-        return getattr(module, attr_name)
+        # If attr_name is None, return the module itself
+        attr = module if attr_name is None else getattr(module, attr_name)
+        # Cache in module __dict__ to avoid repeated __getattr__ calls
+        # and to override the submodule that importlib adds automatically
+        globals()[name] = attr
+        return attr
     raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

tests/unit/test_lazy_imports.py

@@ -59,9 +59,25 @@ def test_lazy_cli_import():
     # CLI module should not be loaded yet
     assert "datajoint.cli" not in sys.modules, "cli module loaded eagerly"
 
-    # Access cli - should trigger lazy load
-    _ = dj.cli
+    # Access cli - should trigger lazy load and return the function
+    cli_func = dj.cli
     assert "datajoint.cli" in sys.modules, "cli module not loaded after access"
+    assert callable(cli_func), "dj.cli should be callable (the cli function)"
+
+
+def test_diagram_module_access():
+    """dj.diagram should return the diagram module for accessing module-level attrs."""
+    # Remove datajoint from sys.modules to get fresh import
+    modules_to_remove = [key for key in sys.modules if key.startswith("datajoint")]
+    for mod in modules_to_remove:
+        del sys.modules[mod]
+
+    import datajoint as dj
+
+    # Access dj.diagram should return the module
+    diagram_module = dj.diagram
+    assert hasattr(diagram_module, "diagram_active"), "diagram module should have diagram_active"
+    assert hasattr(diagram_module, "Diagram"), "diagram module should have Diagram class"
 
 
 def test_diagram_aliases():