reflect PR feedback

Author: DevTGHa

src/designer_plugin/d3sdk/ast_utils.py

@@ -535,8 +535,12 @@ def find_imports_for_function(func: Callable[..., Any]) -> list[PackageInfo]:
                 if not _is_supported_module(alias.name):
                     continue
 
-                # The name used in code is the alias if present, otherwise the module name
-                effective_name = alias.asname if alias.asname else alias.name
+                # The name used in code is the alias if present, otherwise the top-level
+                # package name (e.g. "import logging.handlers" binds "logging", not
+                # "logging.handlers").
+                effective_name = (
+                    alias.asname if alias.asname else alias.name.split(".")[0]
+                )
                 if effective_name in used_names:
                     packages.append(
                         PackageInfo(

src/designer_plugin/d3sdk/client.py

@@ -85,7 +85,7 @@ def create_d3_plugin_method_wrapper(
     2. Serializes the arguments using repr()
     3. Builds a script string in the form: "return plugin.{method_name}({args})"
     4. Creates a PluginPayload with the script and module information
-    5. Sends it to Designer via d3_api_plugin or d3_api_aplugin
+    5. Sends it to Designer via d3_api_execute or d3_api_aexecute
     6. Returns the result from the remote execution
 
     Args:

src/designer_plugin/d3sdk/function.py

@@ -263,15 +263,11 @@ def __init__(self, module_name: str, func: Callable[P, T]):
 
         super().__init__(func)
 
-        # Auto-register packages used by this function
-        D3Function._available_packages[module_name].update(
-            pkg.to_import_statement() for pkg in self._function_info.packages
-        )
-
         # 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]:
+        is_replacement = self in D3Function._available_d3functions[module_name]
+        if is_replacement:
             logger.debug(
                 "Function '%s' in module '%s' is being replaced.",
                 self.name,
@@ -280,6 +276,19 @@ def __init__(self, module_name: str, func: Callable[P, T]):
             D3Function._available_d3functions[module_name].discard(self)
         D3Function._available_d3functions[module_name].add(self)
 
+        if is_replacement:
+            # Full rebuild needed to evict stale imports from the replaced function.
+            D3Function._available_packages[module_name] = {
+                pkg.to_import_statement()
+                for f in D3Function._available_d3functions[module_name]
+                for pkg in f._function_info.packages
+            }
+        else:
+            # New function: incrementally add its packages. No stale imports to remove.
+            D3Function._available_packages[module_name].update(
+                pkg.to_import_statement() for pkg in self._function_info.packages
+            )
+
     def __eq__(self, other: object) -> bool:
         """Check equality based on function name for unique registration.
 

src/designer_plugin/d3sdk/session.py

@@ -35,7 +35,7 @@ def __init__(self, hostname: str, port: int, context_modules: set[str]) -> None:
         Args:
             hostname: The hostname of the Designer instance.
             port: The port number of the Designer instance.
-            context_modules: List of module names to register when entering session context.
+            context_modules: Set of module names to register when entering session context.
         """
         self.hostname: str = hostname
         self.port: int = port

tests/test_ast_utils.py

@@ -5,6 +5,7 @@
 
 import ast
 import inspect
+import logging.handlers
 import textwrap
 import types
 from os.path import join as path_join
@@ -1208,6 +1209,16 @@ def dummy():
         packages = find_imports_for_function(dummy)
         assert packages == []
 
+    def test_dotted_import_effective_name(self):
+        """import logging.handlers binds 'logging' — should match usage of logging.handlers."""
+
+        def uses_logging_handlers():
+            return logging.handlers.RotatingFileHandler("/tmp/x")
+
+        packages = find_imports_for_function(uses_logging_handlers)
+        statements = [p.to_import_statement() for p in packages]
+        assert "import logging.handlers" in statements
+
 
 if __name__ == "__main__":
     pytest.main([__file__, "-v"])

tests/test_core.py

@@ -514,4 +514,44 @@ def func_using_logging():
         assert payload is not None
         assert "import logging" in payload.contents
 
+    def test_new_functions_accumulate_packages(self):
+        """Adding a second function should add its packages without losing the first's."""
+        module = "test_accumulate_pkg_module"
+        D3Function._available_d3functions[module].clear()
+        D3Function._available_packages[module].clear()
+
+        @d3function(module)
+        def func_a():
+            return logging.getLogger("a")
+
+        assert "import logging" in D3Function._available_packages[module]
+
+        @d3function(module)
+        def func_b():
+            import os
+            return os.getcwd()
+
+        # Both packages must be present after adding func_b
+        assert "import logging" in D3Function._available_packages[module]
+        assert "import os" in D3Function._available_packages[module]
+
+    def test_replacement_removes_stale_packages(self):
+        """Replacing a function with one that uses fewer imports should evict stale packages."""
+        module = "test_stale_pkg_module"
+        D3Function._available_d3functions[module].clear()
+        D3Function._available_packages[module].clear()
+
+        @d3function(module)
+        def my_func():  # uses logging
+            return logging.getLogger("x")
+
+        assert "import logging" in D3Function._available_packages[module]
+
+        @d3function(module)
+        def my_func() -> int:  # noqa: F811  # no longer uses logging
+            return 42
+
+        # Stale import from the old version must be gone
+        assert "import logging" not in D3Function._available_packages[module]
+