Switch to lazy module imports

src/forge/actors/__init__.py

@@ -4,6 +4,16 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-from .collector import Collector
+__all__ = ["Policy", "PolicyRouter"]
 
-__all__ = ["Collector"]
+
+def __getattr__(name):
+    if name == "Policy":
+        from .policy import Policy
+
+        return Policy
+    if name == "PolicyRouter":
+        from .policy import PolicyRouter
+
+        return PolicyRouter
+    raise AttributeError("fmodule {__name__} has no attribute {name}")

src/forge/controller/__init__.py

@@ -3,11 +3,6 @@
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
-from .actor import ForgeActor
-from .interface import ServiceInterface, Session, SessionContext
-from .proc_mesh import get_proc_mesh, spawn_actors
-from .service import Service, ServiceConfig
-from .spawn import spawn_service
 
 __all__ = [
     "Service",
@@ -20,3 +15,44 @@
     "get_proc_mesh",
     "ForgeActor",
 ]
+
+
+def __getattr__(name):
+    if name == "Service":
+        from .service import Service
+
+        return Service
+    elif name == "ServiceConfig":
+        from .service import ServiceConfig
+
+        return ServiceConfig
+    elif name == "ServiceInterface":
+        from .interface import ServiceInterface
+
+        return ServiceInterface
+    elif name == "Session":
+        from .interface import Session
+
+        return Session
+    elif name == "SessionContext":
+        from .interface import SessionContext
+
+        return SessionContext
+    elif name == "spawn_service":
+        from .spawn import spawn_service
+
+        return spawn_service
+    elif name == "spawn_actors":
+        from .proc_mesh import spawn_actors
+
+        return spawn_actors
+    elif name == "get_proc_mesh":
+        from .proc_mesh import get_proc_mesh
+
+        return get_proc_mesh
+    elif name == "ForgeActor":
+        from .actor import ForgeActor
+
+        return ForgeActor
+    else:
+        raise AttributeError(f"module {__name__} has no attribute {name}")

src/forge/data/__init__.py

@@ -3,7 +3,18 @@
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
-from .collate import collate_packed
-from .utils import CROSS_ENTROPY_IGNORE_IDX
 
 __all__ = ["collate_packed", "CROSS_ENTROPY_IGNORE_IDX"]
+
+
+def __getattr__(name):
+    if name == "collate_packed":
+        from .collate import collate_packed
+
+        return collate_packed
+    elif name == "CROSS_ENTROPY_IGNORE_IDX":
+        from .utils import CROSS_ENTROPY_IGNORE_IDX
+
+        return CROSS_ENTROPY_IGNORE_IDX
+    else:
+        raise AttributeError(f"module {__name__} has no attribute {name}")