Detect the used python package and add it when register the python automatically

CHANGELOG.md

@@ -11,6 +11,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - **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.
+- **Automatic import detection**: `@d3function` now automatically discovers file-level imports used by the decorated function and includes them in the registered module. In Jupyter notebooks, place imports inside the function body instead.
+
+### Removed
+- `add_packages_in_current_file()`: Removed. Imports are now detected automatically by `@d3function`.
+- `find_packages_in_current_file()`: Removed. Replaced by `find_imports_for_function()`.
 
 ### Changed
 - `d3_api_plugin` has been renamed to `d3_api_execute`.

CONTRIBUTING.md

@@ -31,14 +31,19 @@ Thank you for your interest in contributing to designer-plugin! This document pr
 
 ### Running Tests
 
-Run the full test suite:
+Run unit tests (default):
 ```bash
 uv run pytest
 ```
 
-Run tests with verbose output:
+Run integration tests (requires a running d3 instance):
 ```bash
-uv run pytest -v
+uv run pytest -m integration
+```
+
+Run all tests:
+```bash
+uv run pytest -m ""
 ```
 
 Run specific test file:

README.md

@@ -169,6 +169,14 @@ The Functional API offers two decorators: `@d3pythonscript` and `@d3function`:
   - Functions decorated with the same `module_name` are grouped together and can call each other, enabling function chaining and code reuse.
   - 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"})`).
 
+> **Jupyter Notebook:** File-level imports (e.g., `import numpy as np` in a separate cell) cannot be automatically detected. In Jupyter, place any required imports inside the function body itself:
+> ```python
+> @d3function("mymodule")
+> def my_fn():
+>     import numpy as np
+>     return np.array([1, 2])
+> ```
+
 ### Session API Methods
 
 Both `D3AsyncSession` and `D3Session` provide two methods for executing functions:

pyproject.toml

@@ -109,6 +109,11 @@ python_classes = ["Test*"]
 python_functions = ["test_*"]
 addopts = [
     "-v",
+    "-m", "not integration",
     "--strict-markers",
     "--strict-config",
 ]
+markers = [
+    "integration: tests that require a running d3 instance",
+]
+

src/designer_plugin/d3sdk/__init__.py

@@ -5,7 +5,7 @@
 
 from .client import D3PluginClient
 from .function import (
-    add_packages_in_current_file,
+    PackageInfo,
     d3function,
     d3pythonscript,
     get_all_d3functions,
@@ -18,9 +18,9 @@
     "D3AsyncSession",
     "D3PluginClient",
     "D3Session",
+    "PackageInfo",
     "d3pythonscript",
     "d3function",
-    "add_packages_in_current_file",
     "get_register_payload",
     "get_all_d3functions",
     "get_all_modules",

src/designer_plugin/d3sdk/ast_utils.py

@@ -4,11 +4,68 @@
 """
 
 import ast
+import functools
 import inspect
 import textwrap
 import types
+from collections.abc import Callable
 from typing import Any
 
+from pydantic import BaseModel, Field
+
+from designer_plugin.d3sdk.builtin_modules import SUPPORTED_MODULES
+
+
+###############################################################################
+# Package info models
+class ImportAlias(BaseModel):
+    """Represents a single imported name with an optional alias.
+
+    Mirrors the structure of ast.alias for Pydantic compatibility.
+    """
+
+    name: str = Field(
+        description="The imported name (e.g., 'Path' in 'from pathlib import Path')"
+    )
+    asname: str | None = Field(
+        default=None,
+        description="The alias (e.g., 'np' in 'import numpy as np')",
+    )
+
+
+class PackageInfo(BaseModel):
+    """Structured representation of a Python import statement.
+
+    Rendering rules (via to_import_statement using ast.unparse):
+    - package only              → import package
+    - package + alias           → import package as alias
+    - package + methods         → from package import method1, method2
+    - package + methods w/alias → from package import method1 as alias1
+    """
+
+    package: str = Field(description="The module/package name to import")
+    alias: str | None = Field(
+        default=None,
+        description="Alias for the package (e.g., 'np' in 'import numpy as np')",
+    )
+    methods: list[ImportAlias] = Field(
+        default=[],
+        description="Imported names for 'from X import ...' style imports",
+    )
+
+    def to_import_statement(self) -> str:
+        """Render back to a Python import statement using ast.unparse."""
+        node: ast.stmt
+        if self.methods:
+            node = ast.ImportFrom(
+                module=self.package,
+                names=[ast.alias(name=m.name, asname=m.asname) for m in self.methods],
+                level=0,
+            )
+        else:
+            node = ast.Import(names=[ast.alias(name=self.package, asname=self.alias)])
+        return ast.unparse(node)
+
 
 ###############################################################################
 # Source code extraction utilities
@@ -369,94 +426,138 @@ def validate_and_extract_args(
 
 
 ###############################################################################
-# Python package finder utility
-def find_packages_in_current_file(caller_stack: int = 1) -> list[str]:
-    """Find all import statements in the caller's file by inspecting the call stack.
+# Function-scoped import extraction utility
+def _collect_used_names(func_node: ast.FunctionDef | ast.AsyncFunctionDef) -> set[str]:
+    """Collect all identifier names used inside a function body.
+
+    Walks the function's AST body and extracts:
+    - Simple names (ast.Name nodes, e.g., ``foo`` in ``foo()``)
+    - Root names of attribute chains (e.g., ``np`` in ``np.array()``)
+
+    Args:
+        func_node: The function AST node to analyse.
+
+    Returns:
+        Set of identifier strings used in the function body.
+    """
+    names: set[str] = set()
+    for node in ast.walk(func_node):
+        if isinstance(node, ast.Name):
+            names.add(node.id)
+        elif isinstance(node, ast.Attribute):
+            # Walk down the attribute chain to find the root name
+            root: ast.expr = node
+            while isinstance(root, ast.Attribute):
+                root = root.value
+            if isinstance(root, ast.Name):
+                names.add(root.id)
+    return names
+
+
+def _is_type_checking_block(node: ast.If) -> bool:
+    """Check if an if statement is ``if TYPE_CHECKING:``."""
+    return isinstance(node.test, ast.Name) and node.test.id == "TYPE_CHECKING"
+
+
+def _is_builtin_package(module_name: str) -> bool:
+    """Check if a module name matches python builtin package."""
+    return module_name in SUPPORTED_MODULES
+
+
+@functools.cache
+def _get_module_ast(module: types.ModuleType) -> ast.Module | None:
+    """Return the parsed AST for *module*, cached by module identity."""
+    try:
+        return ast.parse(inspect.getsource(module))
+    except (OSError, TypeError):
+        return None
+
+
+def find_imports_for_function(func: Callable[..., Any]) -> list[PackageInfo]:
+    """Extract import statements used by a function from its source file.
 
-    This function walks up the call stack to find the module where it was called from,
-    then parses that module's source code to extract all import statements that are
-    compatible with Python 2.7 and safe to send to Designer.
+    Inspects the module containing *func*, parses all top-level imports, then
+    filters them down to only those whose imported names are actually referenced
+    inside the function body.
 
     Args:
-        caller_stack: Number of frames to go up the call stack. Default is 1 (immediate caller).
-                     Use higher values to inspect files further up the call chain.
+        func: The callable to analyse.
 
     Returns:
-        Sorted list of unique import statement strings (e.g., "import ast", "from pathlib import Path").
+        Sorted list of :class:`PackageInfo` objects representing the imports
+        used by *func*.
 
     Filters applied:
-        - Excludes imports inside `if TYPE_CHECKING:` blocks (type checking only)
-        - Excludes imports from the 'd3blobgen' package (client-side only)
-        - Excludes imports from the 'typing' module (not supported in Python 2.7)
-        - Excludes imports of this function itself to avoid circular references
+        - Excludes imports inside ``if TYPE_CHECKING:`` blocks
+        - Excludes imports from the ``d3blobgen`` package (client-side only)
+        - Excludes imports from the ``typing`` module (not supported in Python 2.7)
+        - Only includes imports whose names are actually used in the function body
     """
-    # Get the this file frame
-    current_frame: types.FrameType | None = inspect.currentframe()
-    if not current_frame:
+    # --- 1. Get the function's module source ---
+    module = inspect.getmodule(func)
+    if not module:
         return []
 
-    # Get the caller's frame (file where this function is called)
-    caller_frame: types.FrameType | None = current_frame
-    for _ in range(caller_stack):
-        if not caller_frame or not caller_frame.f_back:
-            return []
-        caller_frame = caller_frame.f_back
-
-    if not caller_frame:
+    module_tree = _get_module_ast(module)
+    if module_tree is None:
         return []
 
-    modules: types.ModuleType | None = inspect.getmodule(caller_frame)
-    if not modules:
+    # --- 2. Collect names used inside the function body ---
+    func_source = textwrap.dedent(inspect.getsource(func))
+    func_tree = ast.parse(func_source)
+    if not func_tree.body:
         return []
 
-    source: str = inspect.getsource(modules)
-
-    # Parse the source code
-    tree = ast.parse(source)
-
-    # Get the name of this function to filter it out
-    # For example, we don't want `from core import find_packages_in_current_file`
-    function_name: str = current_frame.f_code.co_name
-    # Skip any package from d3blobgen
-    d3blobgen_package_name: str = "d3blobgen"
-    # typing not supported in python2.7
-    typing_package_name: str = "typing"
+    func_node = func_tree.body[0]
+    if not isinstance(func_node, (ast.FunctionDef, ast.AsyncFunctionDef)):
+        return []
 
-    def is_type_checking_block(node: ast.If) -> bool:
-        """Check if an if statement is 'if TYPE_CHECKING:'"""
-        return isinstance(node.test, ast.Name) and node.test.id == "TYPE_CHECKING"
+    used_names = _collect_used_names(func_node)
 
-    imports: list[str] = []
-    for node in tree.body:
-        # Skip TYPE_CHECKING blocks entirely
-        if isinstance(node, ast.If) and is_type_checking_block(node):
+    # --- 3. Parse file-level imports and filter to used ones ---
+    packages: list[PackageInfo] = []
+    for node in module_tree.body:
+        # Skip TYPE_CHECKING blocks
+        if isinstance(node, ast.If) and _is_type_checking_block(node):
             continue
 
         if isinstance(node, ast.Import):
-            imported_modules: list[str] = [alias.name for alias in node.names]
-            # Skip imports that include d3blobgen
-            if any(d3blobgen_package_name in module for module in imported_modules):
-                continue
-            if any(typing_package_name in module for module in imported_modules):
-                continue
-            import_text: str = f"import {', '.join(imported_modules)}"
-            imports.append(import_text)
+            for alias in node.names:
+                if not _is_builtin_package(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
+                if effective_name in used_names:
+                    packages.append(
+                        PackageInfo(
+                            package=alias.name,
+                            alias=alias.asname,
+                        )
+                    )
 
         elif isinstance(node, ast.ImportFrom):
-            imported_module: str | None = node.module
-            imported_names: list[str] = [alias.name for alias in node.names]
-            if not imported_module:
-                continue
-            # Skip imports that include d3blobgen
-            if d3blobgen_package_name in imported_module:
+            if not node.module:
                 continue
-            elif typing_package_name in imported_module:
-                continue
-            # Skip imports that include this function itself
-            if function_name in imported_names:
+            if not _is_builtin_package(node.module):
                 continue
 
-            line_text = f"from {imported_module} import {', '.join(imported_names)}"
-            imports.append(line_text)
+            # Filter to only methods actually used by the function
+            matched_methods: list[ImportAlias] = []
+            for alias in node.names:
+                effective_name = alias.asname if alias.asname else alias.name
+                if effective_name in used_names:
+                    matched_methods.append(
+                        ImportAlias(name=alias.name, asname=alias.asname)
+                    )
+
+            if matched_methods:
+                packages.append(
+                    PackageInfo(
+                        package=node.module,
+                        methods=matched_methods,
+                    )
+                )
 
-    return sorted(set(imports))
+    # Sort by import statement string for deterministic output
+    return sorted(packages, key=lambda p: p.to_import_statement())