⚡️ Speed up function _map_usage by 47%

[codeflash-ai]

📄 47% (0.47x) speedup for _map_usage in pydantic_ai_slim/pydantic_ai/models/mistral.py

⏱️ Runtime : 254 microseconds → 173 microseconds (best of 149 runs)

📝 Explanation and details

REFINEMENT Here is the rewritten program with improved runtime efficiency while preserving existing comments and function names.
The key optimization here is to avoid attribute lookups multiple times by using a local variable and directly initializing Usage via __init__ with positional arguments for faster execution.

Explanation of changes for efficiency:

• Assigned response.usage to a local variable to avoid repeated attribute lookups.
• Used positional arguments for Usage instantiation, which is slightly faster than using keyword arguments in CPython.

The function and returned results are unchanged, but performance is minimally improved, especially in tight loops or high call volumes.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

from types import SimpleNamespace

# imports
import pytest  # used for our unit tests
from pydantic_ai.models.mistral import _map_usage

# function to test
try:
    from mistralai import CompletionChunk as MistralCompletionChunk
    from mistralai.models import \
        ChatCompletionResponse as MistralChatCompletionResponse
except ImportError as e:  # pragma: lax no cover
    # For testability, define stand-ins if mistral is not installed
    class MistralCompletionChunk(SimpleNamespace):
        pass
    class MistralChatCompletionResponse(SimpleNamespace):
        pass

# Minimal stand-in for Usage for testing
class Usage:
    def __init__(self, request_tokens=0, response_tokens=0, total_tokens=0, details=None):
        self.request_tokens = request_tokens
        self.response_tokens = response_tokens
        self.total_tokens = total_tokens
        self.details = details

    def __eq__(self, other):
        if not isinstance(other, Usage):
            return False
        return (self.request_tokens == other.request_tokens and
                self.response_tokens == other.response_tokens and
                self.total_tokens == other.total_tokens and
                self.details == other.details)

    def __repr__(self):
        return (f"Usage(request_tokens={self.request_tokens}, "
                f"response_tokens={self.response_tokens}, "
                f"total_tokens={self.total_tokens}, details={self.details})")
from pydantic_ai.models.mistral import _map_usage

# unit tests

# --- Basic Test Cases ---


















import pytest  # used for our unit tests
from pydantic_ai.models.mistral import _map_usage
from pydantic_ai.usage import Usage

# function to test
try:
    from mistralai import CompletionChunk as MistralCompletionChunk
    from mistralai.models import \
        ChatCompletionResponse as MistralChatCompletionResponse
except ImportError as e:  # pragma: lax no cover
    raise ImportError(
        'Please install `mistral` to use the Mistral model, '
        'you can use the `mistral` optional group — `pip install "pydantic-ai-slim[mistral]"`'
    ) from e
from pydantic_ai.models.mistral import _map_usage

# --- Mocks for testing (since we cannot import actual mistralai classes in this environment) ---
# These are only for testing purposes and mimic the real classes' interfaces.

class MockUsage:
    def __init__(self, prompt_tokens, completion_tokens, total_tokens):
        self.prompt_tokens = prompt_tokens
        self.completion_tokens = completion_tokens
        self.total_tokens = total_tokens

class MockResponse:
    def __init__(self, usage=None):
        self.usage = usage

# --- Unit tests ---

# 1. Basic Test Cases

def test_map_usage_basic_normal_usage():
    """Test mapping with a typical usage object."""
    usage = MockUsage(prompt_tokens=10, completion_tokens=20, total_tokens=30)
    response = MockResponse(usage=usage)
    codeflash_output = _map_usage(response); result = codeflash_output # 916ns -> 667ns (37.3% faster)

def test_map_usage_zero_tokens():
    """Test mapping with all usage tokens set to zero."""
    usage = MockUsage(prompt_tokens=0, completion_tokens=0, total_tokens=0)
    response = MockResponse(usage=usage)
    codeflash_output = _map_usage(response); result = codeflash_output # 667ns -> 583ns (14.4% faster)

def test_map_usage_none_usage():
    """Test mapping when usage is None."""
    response = MockResponse(usage=None)
    codeflash_output = _map_usage(response); result = codeflash_output # 459ns -> 458ns (0.218% faster)

# 2. Edge Test Cases

def test_map_usage_negative_tokens():
    """Test with negative token counts (should map as-is)."""
    usage = MockUsage(prompt_tokens=-1, completion_tokens=-2, total_tokens=-3)
    response = MockResponse(usage=usage)
    codeflash_output = _map_usage(response); result = codeflash_output # 542ns -> 458ns (18.3% faster)

def test_map_usage_large_token_values():
    """Test with very large token values (near integer limits)."""
    big = 2**63 - 1
    usage = MockUsage(prompt_tokens=big, completion_tokens=big, total_tokens=big)
    response = MockResponse(usage=usage)
    codeflash_output = _map_usage(response); result = codeflash_output # 708ns -> 542ns (30.6% faster)

def test_map_usage_partial_usage_object():
    """Test with a usage object missing one or more attributes."""
    class PartialUsage:
        def __init__(self):
            self.prompt_tokens = 5
            # missing completion_tokens and total_tokens
    usage = PartialUsage()
    response = MockResponse(usage=usage)
    # Should raise AttributeError when trying to access missing attributes
    with pytest.raises(AttributeError):
        _map_usage(response) # 750ns -> 708ns (5.93% faster)

def test_map_usage_usage_is_falsey_but_not_none():
    """Test with a usage object that is falsey but not None (e.g., empty object)."""
    class EmptyUsage:
        def __bool__(self):
            return False
    usage = EmptyUsage()
    response = MockResponse(usage=usage)
    codeflash_output = _map_usage(response); result = codeflash_output # 583ns -> 541ns (7.76% faster)

def test_map_usage_usage_is_truthy_but_fields_are_none():
    """Test with a usage object that exists but all fields are None."""
    class NoneFieldsUsage:
        def __init__(self):
            self.prompt_tokens = None
            self.completion_tokens = None
            self.total_tokens = None
    usage = NoneFieldsUsage()
    response = MockResponse(usage=usage)
    codeflash_output = _map_usage(response); result = codeflash_output # 666ns -> 500ns (33.2% faster)

# 3. Large Scale Test Cases

def test_map_usage_many_responses():
    """Test mapping with a large number of different responses (scalability)."""
    for i in range(1000):
        usage = MockUsage(prompt_tokens=i, completion_tokens=i*2, total_tokens=i*3)
        response = MockResponse(usage=usage)
        codeflash_output = _map_usage(response); result = codeflash_output # 248μs -> 168μs (47.3% faster)

def test_map_usage_large_token_sum():
    """Test mapping with tokens summing to a large value."""
    # This is to check that the function does not overflow or error on large sums
    prompt = 999_999
    completion = 888_888
    total = prompt + completion
    usage = MockUsage(prompt_tokens=prompt, completion_tokens=completion, total_tokens=total)
    response = MockResponse(usage=usage)
    codeflash_output = _map_usage(response); result = codeflash_output # 541ns -> 416ns (30.0% faster)



from pydantic_ai.models.mistral import _map_usage

To edit these changes git checkout codeflash/optimize-_map_usage-mddtbo9p and push.