Fix mypy errors

Author: zastrowm

pyproject.toml

@@ -179,6 +179,7 @@ test-integ = [
 prepare = [
     "hatch fmt --linter",
     "hatch fmt --formatter",
+    "hatch run test-lint",
     "hatch test --all"
 ]
 

src/strands/tools/decorator.py

@@ -51,16 +51,17 @@ def my_tool(param1: str, param2: int = 42) -> dict:
     ParamSpec,
     Type,
     TypeVar,
+    Union,
     cast,
     get_type_hints,
     overload,
-    override,
 )
 
 import docstring_parser
 from pydantic import BaseModel, Field, create_model
+from typing_extensions import override
 
-from strands.types.tools import AgentTool, ToolResult, ToolSpec, ToolUse
+from strands.types.tools import AgentTool, JSONSchema, ToolResult, ToolSpec, ToolUse
 
 # Type for wrapped function
 T = TypeVar("T", bound=Callable[..., Any])
@@ -253,7 +254,7 @@ def validate_input(self, input_data: Dict[str, Any]) -> Dict[str, Any]:
 R = TypeVar("R")  # Return type
 
 
-class DecoratedFunctionTool(Generic[P, R], AgentTool):
+class DecoratedFunctionTool(AgentTool, Generic[P, R]):
     """An AgentTool that wraps a function that was decorated with @tool.
 
     This class adapts Python functions decorated with @tool to the AgentTool interface. It handles both direct
@@ -292,7 +293,7 @@ def __init__(
 
         functools.update_wrapper(wrapper=self, wrapped=self.original_function)
 
-    def __get__(self, instance: Any, obj_type=None) -> "DecoratedFunctionTool[P, R]":
+    def __get__(self, instance: Any, obj_type: Optional[Type] = None) -> "DecoratedFunctionTool[P, R]":
         """Descriptor protocol implementation for proper method binding.
 
         This method enables the decorated function to work correctly when used as a class method.
@@ -396,7 +397,8 @@ def invoke(self, tool: ToolUse, *args: Any, **kwargs: dict[str, Any]) -> ToolRes
             if "agent" in kwargs and "agent" in self._metadata.signature.parameters:
                 validated_input["agent"] = kwargs.get("agent")
 
-            result = self.original_function(**validated_input)
+            # We get "too few arguments here" but because that's because fof the way we're calling it
+            result = self.original_function(**validated_input)  # type: ignore
 
             # FORMAT THE RESULT for Strands Agent
             if isinstance(result, dict) and "status" in result and "content" in result:
@@ -456,8 +458,19 @@ def get_display_properties(self) -> dict[str, str]:
 def tool(__func: Callable[P, R]) -> DecoratedFunctionTool[P, R]: ...
 # Handle @decorator()
 @overload
-def tool(**tool_kwargs: Any) -> Callable[[Callable[P, R]], DecoratedFunctionTool[P, R]]: ...
-def tool(func: Optional[Callable[P, R]] = None, **tool_kwargs: Any) -> DecoratedFunctionTool[P, R]:
+def tool(
+    description: Optional[str] = None,
+    inputSchema: Optional[JSONSchema] = None,
+    name: Optional[str] = None,
+) -> Callable[[Callable[P, R]], DecoratedFunctionTool[P, R]]: ...
+# Suppressing the type error because we want callers to be able to use both `tool` and `tool()` at the
+# call site, but the actual implementation handles that and it's not representable via the type-system
+def tool(  # type: ignore
+    func: Optional[Callable[P, R]] = None,
+    description: Optional[str] = None,
+    inputSchema: Optional[JSONSchema] = None,
+    name: Optional[str] = None,
+) -> Union[DecoratedFunctionTool[P, R], Callable[[Callable[P, R]], DecoratedFunctionTool[P, R]]]:
     """Decorator that transforms a Python function into a Strands tool.
 
     This decorator seamlessly enables a function to be called both as a regular Python function and as a Strands tool.
@@ -472,25 +485,31 @@ def tool(func: Optional[Callable[P, R]] = None, **tool_kwargs: Any) -> Decorated
     4. Formats return values according to the expected Strands tool result format
     5. Provides automatic error handling and reporting
 
+    The decorator can be used in two ways:
+    - As a simple decorator: `@tool`
+    - With parameters: `@tool(name="custom_name", description="Custom description")`
+
     Args:
-        func: The function to decorate.
-        **tool_kwargs: Additional tool specification options to override extracted values.
-            E.g., `name="custom_name", description="Custom description"`.
+        func: The function to decorate. When used as a simple decorator, this is the function being decorated.
+            When used with parameters, this will be None.
+        description: Optional custom description to override the function's docstring.
+        inputSchema: Optional custom JSON schema to override the automatically generated schema.
+        name: Optional custom name to override the function's name.
 
     Returns:
-        The decorated function with attached tool specifications.
+        An AgentTool that also mimics the original function when invoked
 
     Example:
         ```python
         @tool
         def my_tool(name: str, count: int = 1) -> str:
             '''Does something useful with the provided parameters.
 
-            "Args:
+    Args:
                 name: The name to process
                 count: Number of times to process (default: 1)
 
-            "Returns:
+    Returns:
                 A message with the result
             '''
             return f"Processed {name} {count} times"
@@ -503,13 +522,25 @@ def my_tool(name: str, count: int = 1) -> str:
         #   "content": [{"text": "Processed example 3 times"}]
         # }
         ```
+
+    Example with parameters:
+        ```python
+        @tool(name="custom_tool", description="A tool with a custom name and description")
+        def my_tool(name: str, count: int = 1) -> str:
+            return f"Processed {name} {count} times"
+        ```
     """
 
     def decorator(f: T) -> "DecoratedFunctionTool[P, R]":
         # Create function tool metadata
         tool_meta = FunctionToolMetadata(f)
         tool_spec = tool_meta.extract_metadata()
-        tool_spec.update(tool_kwargs)
+        if name is not None:
+            tool_spec["name"] = name
+        if description is not None:
+            tool_spec["description"] = description
+        if inputSchema is not None:
+            tool_spec["inputSchema"] = inputSchema
 
         tool_name = tool_spec.get("name", f.__name__)
 
@@ -520,6 +551,8 @@ def decorator(f: T) -> "DecoratedFunctionTool[P, R]":
 
     # Handle both @tool and @tool() syntax
     if func is None:
+        # Need to ignore type-checking here since it's hard to represent the support
+        # for both flows using the type system
         return decorator
 
     return decorator(func)

src/strands/tools/loader.py

@@ -6,7 +6,7 @@
 import os
 import sys
 from pathlib import Path
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, cast
 
 from ..types.tools import AgentTool
 from .decorator import DecoratedFunctionTool
@@ -141,7 +141,8 @@ def load_python_tool(tool_path: str, tool_name: str) -> AgentTool:
                         logger.debug(
                             "tool_name=<%s>, module_path=<%s> | found function-based tool", function_name, module_path
                         )
-                        return func
+                        # mypy has problems converting between DecoratedFunctionTool <-> AgentTool
+                        return cast(AgentTool, func)
                     else:
                         raise ValueError(
                             f"Function {function_name} in {module_path} is not a valid tool (missing @tool decorator)"
@@ -174,8 +175,8 @@ def load_python_tool(tool_path: str, tool_name: str) -> AgentTool:
                     logger.debug(
                         "tool_name=<%s>, tool_path=<%s> | found function-based tool in path", attr_name, tool_path
                     )
-                    # Return as DecoratedFunctionTool
-                    return attr
+                    # mypy has problems converting between DecoratedFunctionTool <-> AgentTool
+                    return cast(AgentTool, attr)
 
             # If no function-based tools found, fall back to traditional module-level tool
             tool_spec = getattr(module, "TOOL_SPEC", None)

src/strands/tools/registry.py

@@ -11,13 +11,14 @@
 from importlib import import_module, util
 from os.path import expanduser
 from pathlib import Path
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, overload
 
 from typing_extensions import TypedDict, cast
 
 from ..types.tools import AgentTool, Tool, ToolChoice, ToolChoiceAuto, ToolConfig, ToolSpec
+from .decorator import DecoratedFunctionTool
 from .loader import scan_module_for_tools
-from .tools import FunctionTool, PythonAgentTool, normalize_schema, normalize_tool_spec
+from .tools import PythonAgentTool, normalize_schema, normalize_tool_spec
 
 logger = logging.getLogger(__name__)
 
@@ -92,15 +93,7 @@ def process_tools(self, tools: List[Any]) -> List[str]:
                     if not function_tools:
                         logger.warning("tool_name=<%s>, module_path=<%s> | invalid agent tool", tool_name, module_path)
 
-            # Case 5: Function decorated with @tool
-            elif inspect.isfunction(tool) and hasattr(tool, "TOOL_SPEC"):
-                try:
-                    function_tool = FunctionTool(tool)
-                    logger.debug("tool_name=<%s> | registering function tool", function_tool.tool_name)
-                    self.register_tool(function_tool)
-                    tool_names.append(function_tool.tool_name)
-                except Exception as e:
-                    logger.warning("tool_name=<%s> | failed to register function tool | %s", tool.__name__, e)
+            # Case 5: AgentTools (which also covers @tool)
             elif isinstance(tool, AgentTool):
                 self.register_tool(tool)
                 tool_names.append(tool.tool_name)
@@ -176,7 +169,12 @@ def get_all_tools_config(self) -> Dict[str, Any]:
         logger.debug("tool_count=<%s> | tools configured", len(tool_config))
         return tool_config
 
-    def register_tool(self, tool: AgentTool) -> None:
+    # mypy has problems converting between DecoratedFunctionTool <-> AgentTool
+    @overload
+    def register_tool(self, tool: DecoratedFunctionTool) -> None: ...
+    @overload
+    def register_tool(self, tool: AgentTool) -> None: ...
+    def register_tool(self, tool: AgentTool) -> None:  # type: ignore
         """Register a tool function with the given name.
 
         Args:

tests/strands/tools/test_decorator.py

@@ -58,14 +58,7 @@ def test_tool(param1: str, param2: int) -> str:
 
     # Make sure these are set properly
     assert test_tool.__wrapped__ is not None
-    assert test_tool.__doc__ == (
-        "Test tool function.\n"
-        "\n"
-        "        Args:\n"
-        "            param1: First parameter\n"
-        "            param2: Second parameter\n"
-        "        "
-    )
+    assert test_tool.__doc__ == test_tool.original_function.__doc__
 
 
 def test_tool_with_custom_name_description():