Merge pull request #6 from disguise-one/add-lazy-module-register

Add register module in lazy manner

Author: DevTGHa

.github/workflows/ci.yml

@@ -2,16 +2,17 @@ name: CI
 
 on:
   push:
-    branches: [main]
+    branches: [main, dev]
   pull_request:
-    branches: [main]
+    branches: [main, dev]
 
 jobs:
   test:
     runs-on: ubuntu-latest
     strategy:
+      fail-fast: false
       matrix:
-        python-version: ["3.11"]
+        python-version: ["3.11", "3.12", "3.13"]
 
     steps:
       - uses: actions/checkout@v4

CHANGELOG.md

@@ -7,9 +7,14 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [1.3.0] - 2026-01-06
 
+### Added
+- **Lazy module registration**: `D3Session.execute()` and `D3AsyncSession.execute()` now automatically register a `@d3function` module on first use, eliminating the need to declare all modules in `context_modules` upfront.
+- `registered_modules` tracking on session instances prevents duplicate registration calls.
+
 ### Changed
 - `d3_api_plugin` has been renamed to `d3_api_execute`.
 - `d3_api_aplugin` has been renamed to `d3_api_aexecute`.
+- `context_modules` parameter type updated from `list[str]` to `set[str]` on `D3Session`, `D3AsyncSession`, and `D3SessionBase`.
 - Updated documentation to reflect `pystub` proxy support.
 
 ## [1.2.0] - 2025-12-02

README.md

@@ -167,7 +167,7 @@ The Functional API offers two decorators: `@d3pythonscript` and `@d3function`:
 - **`@d3function`**:
   - Must be registered on Designer before execution.
   - Functions decorated with the same `module_name` are grouped together and can call each other, enabling function chaining and code reuse.
-  - Registration is automatic when you pass module names to the session context manager (e.g., `D3AsyncSession('localhost', 80, ["mymodule"])`). If you don't provide module names, no registration occurs.
+  - Registration happens automatically on the first call to `execute()` or `rpc()` that references the module — no need to declare modules upfront. You can also pre-register specific modules by passing them to the session context manager (e.g., `D3AsyncSession('localhost', 80, {"mymodule"})`).
 
 ### Session API Methods
 
@@ -209,11 +209,11 @@ def my_time() -> str:
     return str(datetime.datetime.now())
 
 # Usage with async session
-async with D3AsyncSession('localhost', 80, ["mymodule"]) as session:
+async with D3AsyncSession('localhost', 80) as session:
     # d3pythonscript: no registration needed
     await session.rpc(rename_surface.payload("surface 1", "surface 2"))
 
-    # d3function: registered automatically via context manager
+    # d3function: module is registered automatically on first call
     time: str = await session.rpc(
         rename_surface_get_time.payload("surface 1", "surface 2"))
 
@@ -226,7 +226,7 @@ async with D3AsyncSession('localhost', 80, ["mymodule"]) as session:
 
 # Sync usage
 from designer_plugin.d3sdk import D3Session
-with D3Session('localhost', 80, ["mymodule"]) as session:
+with D3Session('localhost', 80) as session:
     session.rpc(rename_surface.payload("surface 1", "surface 2"))
 ```
 

src/designer_plugin/d3sdk/session.py

@@ -29,7 +29,7 @@
 class D3SessionBase:
     """Base class for Designer session management."""
 
-    def __init__(self, hostname: str, port: int, context_modules: list[str]) -> None:
+    def __init__(self, hostname: str, port: int, context_modules: set[str]) -> None:
         """Initialize base session with connection details and module context.
 
         Args:
@@ -39,7 +39,8 @@ def __init__(self, hostname: str, port: int, context_modules: list[str]) -> None
         """
         self.hostname: str = hostname
         self.port: int = port
-        self.context_modules: list[str] = context_modules
+        self.context_modules: set[str] = context_modules
+        self.registered_modules: set[str] = set()
 
 
 class D3Session(D3SessionBase):
@@ -53,7 +54,7 @@ def __init__(
         self,
         hostname: str,
         port: int = D3_PLUGIN_DEFAULT_PORT,
-        context_modules: list[str] | None = None,
+        context_modules: set[str] | None = None,
     ) -> None:
         """Initialize synchronous Designer session.
 
@@ -62,7 +63,7 @@ def __init__(
             port: The port number of the Designer instance.
             context_modules: Optional list of module names to register when entering session context.
         """
-        super().__init__(hostname, port, context_modules or [])
+        super().__init__(hostname, port, context_modules or set())
 
     def __enter__(self) -> "D3Session":
         """Enter context manager and register all context modules.
@@ -117,6 +118,9 @@ def execute(
         Raises:
             PluginException: If the plugin execution fails.
         """
+        if payload.moduleName and payload.moduleName not in self.registered_modules:
+            self.register_module(payload.moduleName)
+
         return d3_api_execute(self.hostname, self.port, payload, timeout_sec)
 
     def request(self, method: Method, url_endpoint: str, **kwargs: Any) -> Any:
@@ -152,6 +156,7 @@ def register_module(
         )
         if payload:
             d3_api_register_module(self.hostname, self.port, payload, timeout_sec)
+            self.registered_modules.add(module_name)
             return True
         return False
 
@@ -186,7 +191,7 @@ def __init__(
         self,
         hostname: str,
         port: int = D3_PLUGIN_DEFAULT_PORT,
-        context_modules: list[str] | None = None,
+        context_modules: set[str] | None = None,
     ) -> None:
         """Initialize asynchronous Designer session.
 
@@ -195,7 +200,7 @@ def __init__(
             port: The port number of the Designer instance.
             context_modules: Optional list of module names to register when entering session context.
         """
-        super().__init__(hostname, port, context_modules or [])
+        super().__init__(hostname, port, context_modules or set())
 
     async def __aenter__(self) -> "D3AsyncSession":
         """Enter async context manager and register all context modules.
@@ -270,6 +275,9 @@ async def execute(
         Raises:
             PluginException: If the plugin execution fails.
         """
+        if payload.moduleName and payload.moduleName not in self.registered_modules:
+            await self.register_module(payload.moduleName)
+
         return await d3_api_aexecute(self.hostname, self.port, payload, timeout_sec)
 
     async def register_module(
@@ -294,6 +302,7 @@ async def register_module(
             await d3_api_aregister_module(
                 self.hostname, self.port, payload, timeout_sec
             )
+            self.registered_modules.add(module_name)
             return True
         return False
 

tests/test_session.py

@@ -0,0 +1,172 @@
+"""
+MIT License
+Copyright (c) 2025 Disguise Technologies ltd
+"""
+
+import asyncio
+from unittest.mock import AsyncMock, patch
+
+from designer_plugin.d3sdk.function import d3function
+from designer_plugin.d3sdk.session import D3AsyncSession, D3Session
+from designer_plugin.models import PluginPayload, PluginResponse, PluginStatus
+
+
+# Register a module so D3Function._available_d3functions knows about it.
+@d3function("lazy_test_module")
+def _lazy_test_fn() -> str:
+    return "hello world"
+
+
+def _make_response() -> PluginResponse:
+    return PluginResponse(
+        status=PluginStatus(code=0, message="OK", details=[]),
+        returnValue=None,
+    )
+
+
+def _module_payload() -> PluginPayload:
+    """Payload that references a registered @d3function module."""
+    return PluginPayload(moduleName="lazy_test_module", script="return _lazy_test_fn()")
+
+
+def _script_payload() -> PluginPayload:
+    """Payload with no module (equivalent to @d3pythonscript)."""
+    return PluginPayload(moduleName=None, script="return 42")
+
+
+class TestD3SessionLazyRegistration:
+    """Lazy registration behaviour for the synchronous D3Session."""
+
+    def test_registered_modules_starts_empty(self):
+        session = D3Session("localhost", 80)
+        assert session.registered_modules == set()
+
+    def test_module_registered_on_first_execute(self):
+        session = D3Session("localhost", 80)
+        with (
+            patch("designer_plugin.d3sdk.session.d3_api_register_module") as mock_reg,
+            patch("designer_plugin.d3sdk.session.d3_api_execute", return_value=_make_response()),
+        ):
+            session.execute(_module_payload())
+            mock_reg.assert_called_once()
+            assert "lazy_test_module" in session.registered_modules
+
+    def test_module_not_re_registered_on_second_execute(self):
+        session = D3Session("localhost", 80)
+        with (
+            patch("designer_plugin.d3sdk.session.d3_api_register_module") as mock_reg,
+            patch("designer_plugin.d3sdk.session.d3_api_execute", return_value=_make_response()),
+        ):
+            session.execute(_module_payload())
+            session.execute(_module_payload())
+            mock_reg.assert_called_once()
+
+    def test_no_registration_for_script_payload(self):
+        """Payloads without a moduleName must never trigger registration."""
+        session = D3Session("localhost", 80)
+        with (
+            patch("designer_plugin.d3sdk.session.d3_api_register_module") as mock_reg,
+            patch("designer_plugin.d3sdk.session.d3_api_execute", return_value=_make_response()),
+        ):
+            session.execute(_script_payload())
+            mock_reg.assert_not_called()
+
+    def test_context_module_not_re_registered_lazily(self):
+        """A module pre-registered via context_modules must not be registered again in execute()."""
+        with (
+            patch("designer_plugin.d3sdk.session.d3_api_register_module") as mock_reg,
+            patch("designer_plugin.d3sdk.session.d3_api_execute", return_value=_make_response()),
+        ):
+            with D3Session("localhost", 80, {"lazy_test_module"}) as session:
+                assert "lazy_test_module" in session.registered_modules
+                session.execute(_module_payload())
+            mock_reg.assert_called_once()  # only from __enter__, not from execute()
+
+    def test_registered_modules_updated_after_execute(self):
+        session = D3Session("localhost", 80)
+        assert "lazy_test_module" not in session.registered_modules
+        with (
+            patch("designer_plugin.d3sdk.session.d3_api_register_module"),
+            patch("designer_plugin.d3sdk.session.d3_api_execute", return_value=_make_response()),
+        ):
+            session.execute(_module_payload())
+        assert "lazy_test_module" in session.registered_modules
+
+
+class TestD3AsyncSessionLazyRegistration:
+    """Lazy registration behaviour for the asynchronous D3AsyncSession."""
+
+    def test_registered_modules_starts_empty(self):
+        session = D3AsyncSession("localhost", 80)
+        assert session.registered_modules == set()
+
+    def test_module_registered_on_first_execute(self):
+        async def run():
+            session = D3AsyncSession("localhost", 80)
+            with (
+                patch("designer_plugin.d3sdk.session.d3_api_aregister_module", new_callable=AsyncMock) as mock_reg,
+                patch("designer_plugin.d3sdk.session.d3_api_aexecute", new_callable=AsyncMock) as mock_exec,
+            ):
+                mock_exec.return_value = _make_response()
+                await session.execute(_module_payload())
+                mock_reg.assert_called_once()
+                assert "lazy_test_module" in session.registered_modules
+
+        asyncio.run(run())
+
+    def test_module_not_re_registered_on_second_execute(self):
+        async def run():
+            session = D3AsyncSession("localhost", 80)
+            with (
+                patch("designer_plugin.d3sdk.session.d3_api_aregister_module", new_callable=AsyncMock) as mock_reg,
+                patch("designer_plugin.d3sdk.session.d3_api_aexecute", new_callable=AsyncMock) as mock_exec,
+            ):
+                mock_exec.return_value = _make_response()
+                await session.execute(_module_payload())
+                await session.execute(_module_payload())
+                mock_reg.assert_called_once()
+
+        asyncio.run(run())
+
+    def test_no_registration_for_script_payload(self):
+        """Payloads without a moduleName must never trigger registration."""
+        async def run():
+            session = D3AsyncSession("localhost", 80)
+            with (
+                patch("designer_plugin.d3sdk.session.d3_api_aregister_module", new_callable=AsyncMock) as mock_reg,
+                patch("designer_plugin.d3sdk.session.d3_api_aexecute", new_callable=AsyncMock) as mock_exec,
+            ):
+                mock_exec.return_value = _make_response()
+                await session.execute(_script_payload())
+                mock_reg.assert_not_called()
+
+        asyncio.run(run())
+
+    def test_context_module_not_re_registered_lazily(self):
+        """A module pre-registered via context_modules must not be registered again in execute()."""
+        async def run():
+            with (
+                patch("designer_plugin.d3sdk.session.d3_api_aregister_module", new_callable=AsyncMock) as mock_reg,
+                patch("designer_plugin.d3sdk.session.d3_api_aexecute", new_callable=AsyncMock) as mock_exec,
+            ):
+                mock_exec.return_value = _make_response()
+                async with D3AsyncSession("localhost", 80, {"lazy_test_module"}) as session:
+                    assert "lazy_test_module" in session.registered_modules
+                    await session.execute(_module_payload())
+                mock_reg.assert_called_once()
+
+        asyncio.run(run())
+
+    def test_registered_modules_updated_after_execute(self):
+        async def run():
+            session = D3AsyncSession("localhost", 80)
+            assert "lazy_test_module" not in session.registered_modules
+            with (
+                patch("designer_plugin.d3sdk.session.d3_api_aregister_module", new_callable=AsyncMock),
+                patch("designer_plugin.d3sdk.session.d3_api_aexecute", new_callable=AsyncMock) as mock_exec,
+            ):
+                mock_exec.return_value = _make_response()
+                await session.execute(_module_payload())
+            assert "lazy_test_module" in session.registered_modules
+
+        asyncio.run(run())