Merge pull request #7 from disguise-one/jupyter-support

Jupyter support

Author: DevTGHa

.github/workflows/ci.yml

@@ -15,10 +15,10 @@ jobs:
         python-version: ["3.11", "3.12", "3.13"]
 
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v6
 
       - name: Install uv
-        uses: astral-sh/setup-uv@v5
+        uses: astral-sh/setup-uv@v7
         with:
           enable-cache: true
 

.github/workflows/test-publish.yml

@@ -0,0 +1,64 @@
+name: Publish to Test PyPI
+
+on:
+  workflow_dispatch:
+
+permissions:
+  contents: read
+  id-token: write  # Required for trusted publishing
+
+jobs:
+  test-publish:
+    runs-on: ubuntu-latest
+
+    steps:
+    - name: Checkout repository
+      uses: actions/checkout@v6
+
+    - name: Set up Python
+      uses: actions/setup-python@v6
+      with:
+        python-version: '3.12'  # tomllib requires >= 3.11
+
+    - name: Install uv
+      uses: astral-sh/setup-uv@v7
+
+    - name: Append .dev suffix for unique Test PyPI versions
+      run: |
+        python -c "
+        import tomllib, pathlib, re
+        path = pathlib.Path('pyproject.toml')
+        text = path.read_text()
+        data = tomllib.loads(text)
+        version = data['project']['version']
+        dev_version = f'{version}.dev${{ github.run_number }}'
+        # Only replace the version inside the [project] section to avoid
+        # accidentally matching a version key in [tool.*] sections.
+        def replace_in_project_section(text, old_ver, new_ver):
+            project_match = re.search(r'^\[project\]', text, re.MULTILINE)
+            if not project_match:
+                raise RuntimeError('[project] section not found in pyproject.toml')
+            start = project_match.start()
+            # Find the next top-level section header or end of file
+            next_section = re.search(r'^\[(?!project[.\]])', text[start+1:], re.MULTILINE)
+            end = (start + 1 + next_section.start()) if next_section else len(text)
+            section = text[start:end]
+            section = re.sub(
+                r'(version\s*=\s*\")' + re.escape(old_ver) + r'\"',
+                r'\g<1>' + new_ver + '\"',
+                section, count=1,
+            )
+            return text[:start] + section + text[end:]
+        text = replace_in_project_section(text, version, dev_version)
+        path.write_text(text)
+        print(f'Version set to {dev_version}')
+        "
+
+    - name: Build package
+      run: uv build
+
+    - name: Publish to Test PyPI
+      uses: pypa/gh-action-pypi-publish@release/v1
+      with:
+        repository-url: https://test.pypi.org/legacy/
+        skip-existing: true

CHANGELOG.md

@@ -10,12 +10,15 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### 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.
+- **Jupyter notebook support**: `@d3function` now automatically replaces a previously registered function when the same name is re-registered in the same module, with a warning log. This enables iterative workflows in Jupyter notebooks where cells are re-executed.
 
 ### 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.
+- Bumped `actions/checkout` to v6 and `astral-sh/setup-uv` to v7 in CI.
+- Added Test PyPI publish workflow (`test-publish.yml`) for dev version releases.
 
 ## [1.2.0] - 2025-12-02
 

src/designer_plugin/d3sdk/function.py

@@ -6,6 +6,7 @@
 import ast
 import functools
 import inspect
+import logging
 import textwrap
 from collections import defaultdict
 from collections.abc import Callable
@@ -24,6 +25,8 @@
     RegisterPayload,
 )
 
+logger = logging.getLogger(__name__)
+
 
 ###############################################################################
 # Plugin function related implementations
@@ -253,6 +256,16 @@ def __init__(self, module_name: str, func: Callable[P, T]):
 
         super().__init__(func)
 
+        # Update the function in case the function was updated in the same session.
+        # For example, jupyter notebook server can be running, but function signature can
+        # change constantly.
+        if self in D3Function._available_d3functions[module_name]:
+            logger.debug(
+                "Function '%s' in module '%s' is being replaced.",
+                self.name,
+                module_name,
+            )
+            D3Function._available_d3functions[module_name].discard(self)
         D3Function._available_d3functions[module_name].add(self)
 
     def __eq__(self, other: object) -> bool:

src/designer_plugin/d3sdk/session.py

@@ -61,7 +61,7 @@ def __init__(
         Args:
             hostname: The hostname of the Designer instance.
             port: The port number of the Designer instance.
-            context_modules: Optional list of module names to register when entering session context.
+            context_modules: Optional set of module names to register when entering session context.
         """
         super().__init__(hostname, port, context_modules or set())
 
@@ -198,7 +198,7 @@ def __init__(
         Args:
             hostname: The hostname of the Designer instance.
             port: The port number of the Designer instance.
-            context_modules: Optional list of module names to register when entering session context.
+            context_modules: Optional set of module names to register when entering session context.
         """
         super().__init__(hostname, port, context_modules or set())
 

tests/test_core.py

@@ -3,6 +3,8 @@
 Copyright (c) 2025 Disguise Technologies ltd
 """
 
+import logging
+
 import pytest
 
 from designer_plugin.d3sdk.function import (
@@ -329,6 +331,61 @@ def test_inequality_different_functions(self):
 
 
 
+class TestD3FunctionReplacement:
+    """Test that re-registering a D3Function with the same name replaces the old one."""
+
+    def test_reregister_replaces_function(self):
+        """Re-registering a function with the same name should replace it in the set."""
+        module = "test_replace_module"
+        D3Function._available_d3functions[module].clear()
+
+        @d3function(module)
+        def my_func(a: int) -> int:
+            return a
+
+        @d3function(module)
+        def my_func(a: int, b: int) -> int:  # noqa: F811
+            return a + b
+
+        funcs = D3Function._available_d3functions[module]
+        matching = [f for f in funcs if f.name == "my_func"]
+        assert len(matching) == 1
+        assert matching[0].function_info.args == ["a", "b"]
+
+    def test_reregister_logs_debug(self, caplog):
+        """Re-registering should log a debug message."""
+        module = "test_replace_warn_module"
+        D3Function._available_d3functions[module].clear()
+
+        @d3function(module)
+        def warn_func() -> None:
+            pass
+
+        with caplog.at_level(logging.DEBUG, logger="designer_plugin.d3sdk.function"):
+            @d3function(module)
+            def warn_func() -> int:  # noqa: F811
+                return 1
+
+        assert any("warn_func" in msg and "being replaced" in msg for msg in caplog.messages)
+
+    def test_set_size_unchanged_after_replacement(self):
+        """The function set size should stay the same after replacement."""
+        module = "test_replace_size_module"
+        D3Function._available_d3functions[module].clear()
+
+        @d3function(module)
+        def size_func(x: int) -> int:
+            return x
+
+        assert len(D3Function._available_d3functions[module]) == 1
+
+        @d3function(module)
+        def size_func(x: int, y: int) -> int:  # noqa: F811
+            return x + y
+
+        assert len(D3Function._available_d3functions[module]) == 1
+
+
 class TestD3PythonScript:
     def test_d3pythonscript_decorator(self):
         @d3pythonscript