⚡️ Speed up function _customize_tool_def by 30%

[codeflash-ai]

📄 30% (0.30x) speedup for _customize_tool_def in pydantic_ai_slim/pydantic_ai/models/__init__.py

⏱️ Runtime : 77.1 microseconds → 59.5 microseconds (best of 329 runs)

📝 Explanation and details

REFINEMENT Let's analyze the line profiler stats.

• The slowest lines are the three replace() calls and the transformer(...).walk() call.
• replace() is called potentially twice per function call, creating new altered copies of ToolDefinition. Allocating new dataclasses is expensive compared to updating in-place. However, since you require equivalent return values, and since dataclasses are typically frozen (frozen=True) for immutability (as it seems here), we can't mutate them directly.
• The walk() method and is_strict_compatible property are called on the transformer. It would help to avoid constructing multiple dataclasses when we could build everything in a single replace call.
• If t.strict is not None, the first replace is skipped. But if it is None, there are two allocations of the dataclass due to two replace() calls.

Optimization plan

• Compute both parameters_json_schema and new strict value before any replace call, and issue only one call to replace().
• Minor: Avoid temporary variables if they aren’t needed.

Key:
If t.strict is None, we want both new parameters_json_schema and an updated strict.
If t.strict is not None, just update parameters_json_schema (leave strict alone).

Here’s the optimized code.

Why is this faster?

• Allocates only one new dataclass per call (was possibly two).
• Walks and computes all needed values up front; no unnecessary extra replace calls.
• Logic is clearer; less overhead from repeated field copying.

---

If you're allowed to mutate t in-place (i.e., if ToolDefinition is NOT frozen or you control its mutability), it could be even faster, but this would change semantics.
This version preserves the semantics, and is the most optimal without change of return values or mutability.

Let me know if mutability is permitted, and I can make this even faster!

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 30 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	✅ 1 Passed
📊 Tests Coverage	100.0%

🌀 Generated Regression Tests and Runtime

from dataclasses import dataclass, replace
from typing import Any, Dict, Optional

# imports
import pytest
from pydantic_ai.models.__init__ import _customize_tool_def

# --- Mocks and minimal implementations for testing ---

@dataclass(frozen=True)
class ToolDefinition:
    parameters_json_schema: Dict[str, Any]
    strict: Optional[bool] = None

# Minimal mock of JsonSchemaTransformer for testing
class JsonSchemaTransformer:
    def __init__(self, schema: Dict[str, Any], strict: Optional[bool] = None):
        self.schema = schema
        self.strict = strict
        # For testability, we allow the schema to define a transformation
        self._walk_result = schema.get("_walk_result", schema)
        self.is_strict_compatible = schema.get("_is_strict_compatible", True)
    def walk(self):
        return self._walk_result
from pydantic_ai.models.__init__ import _customize_tool_def

# --- Unit Tests ---

# 1. Basic Test Cases

def test_basic_no_transformation():
    # Test that if the transformer does not change the schema, output is the same (except strict if None)
    schema = {"type": "object", "properties": {"x": {"type": "integer"}}}
    t = ToolDefinition(parameters_json_schema=schema, strict=True)
    codeflash_output = _customize_tool_def(JsonSchemaTransformer, t); result = codeflash_output # 1.54μs -> 1.54μs (0.000% faster)

def test_basic_strict_none_sets_from_transformer():
    # If strict is None, should use is_strict_compatible from transformer
    schema = {"type": "object", "_is_strict_compatible": False}
    t = ToolDefinition(parameters_json_schema=schema, strict=None)
    codeflash_output = _customize_tool_def(JsonSchemaTransformer, t); result = codeflash_output # 2.21μs -> 1.38μs (60.6% faster)

def test_basic_schema_transformation():
    # If transformer changes schema, output should reflect that
    schema = {"type": "object", "_walk_result": {"type": "array"}, "_is_strict_compatible": True}
    t = ToolDefinition(parameters_json_schema=schema, strict=None)
    codeflash_output = _customize_tool_def(JsonSchemaTransformer, t); result = codeflash_output # 2.12μs -> 1.33μs (59.4% faster)

# 2. Edge Test Cases

def test_edge_empty_schema():
    # Should handle empty schema dict
    schema = {}
    t = ToolDefinition(parameters_json_schema=schema, strict=None)
    codeflash_output = _customize_tool_def(JsonSchemaTransformer, t); result = codeflash_output # 2.08μs -> 1.33μs (56.1% faster)

def test_edge_schema_with_extra_fields():
    # Should not drop extra fields in schema
    schema = {"type": "object", "title": "My Tool", "_walk_result": {"type": "object", "title": "Changed"}}
    t = ToolDefinition(parameters_json_schema=schema, strict=True)
    codeflash_output = _customize_tool_def(JsonSchemaTransformer, t); result = codeflash_output # 1.42μs -> 1.42μs (0.000% faster)

def test_edge_strict_false():
    # If strict is False, should remain False even if _is_strict_compatible is True
    schema = {"type": "object", "_is_strict_compatible": True}
    t = ToolDefinition(parameters_json_schema=schema, strict=False)
    codeflash_output = _customize_tool_def(JsonSchemaTransformer, t); result = codeflash_output # 1.50μs -> 1.42μs (5.86% faster)

def test_edge_strict_true_with_transformer_false():
    # If strict is True, but transformer would set to False, should remain True
    schema = {"type": "object", "_is_strict_compatible": False}
    t = ToolDefinition(parameters_json_schema=schema, strict=True)
    codeflash_output = _customize_tool_def(JsonSchemaTransformer, t); result = codeflash_output # 1.42μs -> 1.38μs (3.05% faster)

def test_edge_schema_with_nonstandard_types():
    # Should handle schemas with nonstandard types (e.g., enums, custom fields)
    schema = {
        "type": "object",
        "properties": {"color": {"type": "string", "enum": ["red", "blue"]}},
        "_walk_result": {"type": "object", "properties": {"color": {"type": "string", "enum": ["red"]}}}
    }
    t = ToolDefinition(parameters_json_schema=schema, strict=None)
    codeflash_output = _customize_tool_def(JsonSchemaTransformer, t); result = codeflash_output # 2.04μs -> 1.38μs (48.5% faster)

def test_edge_schema_with_nested_walk_result():
    # _walk_result can be nested/complex
    schema = {
        "type": "object",
        "properties": {"x": {"type": "object", "properties": {"y": {"type": "integer"}}}},
        "_walk_result": {"type": "object", "properties": {"x": {"type": "object", "properties": {"y": {"type": "string"}}}}}
    }
    t = ToolDefinition(parameters_json_schema=schema, strict=None)
    codeflash_output = _customize_tool_def(JsonSchemaTransformer, t); result = codeflash_output # 2.08μs -> 1.33μs (56.3% faster)

# 3. Large Scale Test Cases

def test_large_schema_many_properties():
    # Test with a schema with many properties (scalability)
    props = {f"key{i}": {"type": "integer"} for i in range(500)}
    schema = {"type": "object", "properties": props}
    t = ToolDefinition(parameters_json_schema=schema, strict=None)
    codeflash_output = _customize_tool_def(JsonSchemaTransformer, t); result = codeflash_output # 2.25μs -> 1.42μs (58.9% faster)

def test_large_schema_transformation_many_properties():
    # Transformer changes all types to string in a large schema
    class StringifyTransformer(JsonSchemaTransformer):
        def walk(self):
            # Change all property types to string
            schema = dict(self.schema)
            if "properties" in schema:
                schema["properties"] = {k: {"type": "string"} for k in schema["properties"]}
            return schema
    props = {f"key{i}": {"type": "integer"} for i in range(300)}
    schema = {"type": "object", "properties": props}
    t = ToolDefinition(parameters_json_schema=schema, strict=None)
    codeflash_output = _customize_tool_def(StringifyTransformer, t); result = codeflash_output # 15.0μs -> 13.6μs (10.7% faster)
    for v in result.parameters_json_schema["properties"].values():
        pass

def test_large_schema_nested_objects():
    # Deeply nested schema (depth 10)
    schema = {"type": "object"}
    current = schema
    for i in range(10):
        current["properties"] = {"nested": {"type": "object"}}
        current = current["properties"]["nested"]
    t = ToolDefinition(parameters_json_schema=schema, strict=None)
    codeflash_output = _customize_tool_def(JsonSchemaTransformer, t); result = codeflash_output # 2.12μs -> 1.33μs (59.4% faster)
    # Check that the nesting is preserved
    cur = result.parameters_json_schema
    for i in range(10):
        cur = cur["properties"]["nested"]

def test_large_schema_with_walk_result():
    # Large schema with _walk_result as a large dict
    walk_result = {"type": "object", "properties": {f"x{i}": {"type": "string"} for i in range(700)}}
    schema = {"type": "object", "_walk_result": walk_result}
    t = ToolDefinition(parameters_json_schema=schema, strict=None)
    codeflash_output = _customize_tool_def(JsonSchemaTransformer, t); result = codeflash_output # 2.17μs -> 1.42μs (52.9% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

from dataclasses import dataclass, replace
from typing import Any, Dict, Optional

# imports
import pytest  # used for our unit tests
from pydantic_ai.models.__init__ import _customize_tool_def


# Mock classes to simulate the real dependencies
@dataclass(frozen=True)
class ToolDefinition:
    parameters_json_schema: Dict[str, Any]
    strict: Optional[bool] = None

class JsonSchemaTransformer:
    def __init__(self, schema: Dict[str, Any], strict: Optional[bool] = None):
        # Save the schema and strictness for use in walk()
        self.schema = schema
        self.strict = strict
        # For testing, let's say strict compatibility is True if schema has "strict_mode": True
        self.is_strict_compatible = schema.get("strict_mode", False)

    def walk(self):
        # Simulate a transformation: add a marker to the schema
        transformed = dict(self.schema)
        transformed["_transformed"] = True
        return transformed
from pydantic_ai.models.__init__ import _customize_tool_def

# unit tests

# ------------------ Basic Test Cases ------------------

def test_basic_transformation_and_replace():
    # Test that the schema is transformed and replaced in the ToolDefinition
    orig_schema = {"type": "object", "properties": {"x": {"type": "integer"}}}
    t = ToolDefinition(parameters_json_schema=orig_schema, strict=True)
    codeflash_output = _customize_tool_def(JsonSchemaTransformer, t); result = codeflash_output # 1.75μs -> 1.67μs (4.98% faster)

def test_basic_strict_none_sets_from_transformer():
    # If strict=None, it should be set from the transformer's is_strict_compatible
    orig_schema = {"type": "object", "strict_mode": True}
    t = ToolDefinition(parameters_json_schema=orig_schema, strict=None)
    codeflash_output = _customize_tool_def(JsonSchemaTransformer, t); result = codeflash_output # 2.38μs -> 1.50μs (58.3% faster)

def test_basic_strict_none_sets_false():
    # If strict=None and schema doesn't have strict_mode, should be False
    orig_schema = {"type": "object"}
    t = ToolDefinition(parameters_json_schema=orig_schema, strict=None)
    codeflash_output = _customize_tool_def(JsonSchemaTransformer, t); result = codeflash_output # 2.21μs -> 1.42μs (55.8% faster)

def test_basic_no_properties():
    # Minimal schema, no properties
    orig_schema = {}
    t = ToolDefinition(parameters_json_schema=orig_schema, strict=True)
    codeflash_output = _customize_tool_def(JsonSchemaTransformer, t); result = codeflash_output # 1.46μs -> 1.46μs (0.069% slower)

# ------------------ Edge Test Cases ------------------

def test_edge_empty_schema_and_strict_none():
    # What if schema is empty and strict is None?
    orig_schema = {}
    t = ToolDefinition(parameters_json_schema=orig_schema, strict=None)
    codeflash_output = _customize_tool_def(JsonSchemaTransformer, t); result = codeflash_output # 2.12μs -> 1.38μs (54.5% faster)

def test_edge_schema_with_unusual_types():
    # Schema with non-standard types/values
    orig_schema = {"type": "customType", "strict_mode": True}
    t = ToolDefinition(parameters_json_schema=orig_schema, strict=None)
    codeflash_output = _customize_tool_def(JsonSchemaTransformer, t); result = codeflash_output # 2.08μs -> 1.42μs (47.1% faster)

def test_edge_schema_with_nested_properties():
    # Deeply nested schema
    orig_schema = {
        "type": "object",
        "properties": {
            "a": {"type": "object", "properties": {"b": {"type": "integer"}}}
        },
        "strict_mode": False
    }
    t = ToolDefinition(parameters_json_schema=orig_schema, strict=None)
    codeflash_output = _customize_tool_def(JsonSchemaTransformer, t); result = codeflash_output # 2.21μs -> 1.46μs (51.4% faster)

def test_edge_schema_with_non_bool_strict_mode():
    # strict_mode is present but not a bool
    orig_schema = {"type": "object", "strict_mode": "yes"}
    t = ToolDefinition(parameters_json_schema=orig_schema, strict=None)
    codeflash_output = _customize_tool_def(JsonSchemaTransformer, t); result = codeflash_output # 2.04μs -> 1.33μs (53.2% faster)

def test_edge_strict_already_false():
    # strict is already False, should not be changed
    orig_schema = {"type": "object", "strict_mode": True}
    t = ToolDefinition(parameters_json_schema=orig_schema, strict=False)
    codeflash_output = _customize_tool_def(JsonSchemaTransformer, t); result = codeflash_output # 1.54μs -> 1.50μs (2.73% faster)

def test_edge_strict_already_true():
    # strict is already True, should not be changed
    orig_schema = {"type": "object", "strict_mode": False}
    t = ToolDefinition(parameters_json_schema=orig_schema, strict=True)
    codeflash_output = _customize_tool_def(JsonSchemaTransformer, t); result = codeflash_output # 1.42μs -> 1.46μs (2.88% slower)

def test_edge_schema_with_additional_keys():
    # Schema with extra keys
    orig_schema = {"type": "object", "foo": 123, "strict_mode": False}
    t = ToolDefinition(parameters_json_schema=orig_schema, strict=None)
    codeflash_output = _customize_tool_def(JsonSchemaTransformer, t); result = codeflash_output # 2.08μs -> 1.46μs (42.9% faster)

# ------------------ Large Scale Test Cases ------------------

def test_large_schema_many_properties():
    # Test with a schema with many properties (up to 1000)
    properties = {f"field{i}": {"type": "integer"} for i in range(1000)}
    orig_schema = {"type": "object", "properties": properties, "strict_mode": True}
    t = ToolDefinition(parameters_json_schema=orig_schema, strict=None)
    codeflash_output = _customize_tool_def(JsonSchemaTransformer, t); result = codeflash_output # 2.42μs -> 1.62μs (48.7% faster)

def test_large_schema_with_nested_objects():
    # Test with a schema with nested objects up to a reasonable depth
    nested = {"type": "object", "properties": {"inner": {"type": "string"}}}
    orig_schema = {"type": "object", "properties": {"a": nested, "b": nested}, "strict_mode": False}
    t = ToolDefinition(parameters_json_schema=orig_schema, strict=None)
    codeflash_output = _customize_tool_def(JsonSchemaTransformer, t); result = codeflash_output # 2.08μs -> 1.42μs (47.0% faster)

def test_large_multiple_calls_consistency():
    # Call the function multiple times with different schemas to ensure no state is leaked
    schemas = [
        {"type": "object", "strict_mode": True},
        {"type": "object", "strict_mode": False},
        {"type": "object"},
    ]
    expected_stricts = [True, False, False]
    for schema, expected in zip(schemas, expected_stricts):
        t = ToolDefinition(parameters_json_schema=schema, strict=None)
        codeflash_output = _customize_tool_def(JsonSchemaTransformer, t); result = codeflash_output # 4.46μs -> 2.83μs (57.3% faster)

def test_large_schema_with_various_types():
    # Schema with many types and properties
    properties = {
        "int_field": {"type": "integer"},
        "str_field": {"type": "string"},
        "arr_field": {"type": "array", "items": {"type": "number"}},
        "obj_field": {"type": "object", "properties": {"x": {"type": "boolean"}}},
    }
    orig_schema = {"type": "object", "properties": properties, "strict_mode": True}
    t = ToolDefinition(parameters_json_schema=orig_schema, strict=None)
    codeflash_output = _customize_tool_def(JsonSchemaTransformer, t); result = codeflash_output # 2.12μs -> 1.42μs (50.0% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

from pydantic_ai.models.__init__ import _customize_tool_def
from pydantic_ai.profiles._json_schema import InlineDefsJsonSchemaTransformer
from pydantic_ai.tools import ToolDefinition

def test__customize_tool_def():
    _customize_tool_def(InlineDefsJsonSchemaTransformer, ToolDefinition('', {}, description='', outer_typed_dict_key=None, strict=None))

To edit these changes git checkout codeflash/optimize-_customize_tool_def-mdetgyjy and push.