⚡️ Speed up function _encode_locations by 11%

[codeflash-ai]

📄 11% (0.11x) speedup for _encode_locations in sentry_sdk/metrics.py

⏱️ Runtime : 10.0 milliseconds → 9.00 milliseconds (best of 78 runs)

📝 Explanation and details

The optimization achieves an 11% speedup by eliminating redundant function lookups and method calls within the hot loop. The key changes are:

Function Reference Hoisting: Moving _sanitize_metric_key, _sanitize_unit, and mapping.get lookups outside the loop saves attribute resolution overhead on each iteration (8,000+ times in the profiled case).

Optimized Dictionary Access Pattern: Replacing mapping.setdefault(mri, []).append(loc) with explicit mapping.get() and conditional assignment eliminates the overhead of setdefault() creating new lists when keys already exist. The profile shows this saves ~2ms in the setdefault line alone.

F-string Formatting: Using f-strings instead of .format() provides marginal but measurable string formatting improvements.

Performance Characteristics: The optimization is most effective for workloads with:

• Many duplicate metric keys (13.7% speedup in the "same key" large-scale test)
• Large numbers of locations (9-13% speedup in 1000+ location tests)
• Complex metric key sanitization (7-13% speedup in edge cases with special characters)

The optimizations preserve exact behavior including the mutation of location dictionaries with "type": "location", making this a safe drop-in replacement that scales better with input size.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import json
import re
from functools import partial
from typing import Any, Dict, Iterable, List, Tuple

# imports
import pytest  # used for our unit tests
from sentry_sdk.metrics import _encode_locations

_sanitize_unit = partial(re.compile(r"[^a-zA-Z0-9_]+").sub, "")
_sanitize_metric_key = partial(re.compile(r"[^a-zA-Z0-9_\-.]+").sub, "_")
from sentry_sdk.metrics import _encode_locations

# unit tests

def decode_result(result: bytes) -> dict:
    """Helper to decode the bytes output from _encode_locations."""
    return json.loads(result.decode("utf-8"))

# 1. Basic Test Cases


def test_multiple_locations_different_keys():
    # Basic: multiple locations with different keys
    timestamp = 42
    code_locations = [
        (("counter", "requests", "second"), {"file": "a.py"}),
        (("gauge", "memory", "byte"), {"file": "b.py"}),
    ]
    result = decode_result(_encode_locations(timestamp, code_locations)) # 26.2μs -> 25.1μs (4.26% faster)


def test_empty_code_locations():
    # Basic: empty code_locations should produce empty mapping
    timestamp = 0
    code_locations = []
    result = decode_result(_encode_locations(timestamp, code_locations)) # 14.9μs -> 15.2μs (1.92% slower)

def test_location_dict_not_mutated():
    # Basic: Ensure original dict is mutated (adds 'type': 'location')
    code_locations = [
        (("counter", "requests", "second"), {"file": "main.py"})
    ]
    loc_dict = code_locations[0][1]
    _encode_locations(0, code_locations) # 19.1μs -> 18.0μs (5.88% faster)

# 2. Edge Test Cases

def test_key_sanitization_special_chars():
    # Edge: metric name and unit with special chars should be sanitized
    code_locations = [
        (("counter", "req!uests", "sec@ond$"), {"file": "main.py"})
    ]
    result = decode_result(_encode_locations(0, code_locations)) # 18.7μs -> 17.5μs (7.15% faster)

def test_metric_name_with_dash_and_dot():
    # Edge: metric name with dash and dot should be preserved
    code_locations = [
        (("gauge", "foo-bar.baz", "unit"), {"file": "x.py"})
    ]
    result = decode_result(_encode_locations(0, code_locations)) # 16.4μs -> 16.1μs (1.93% faster)

def test_unit_with_underscore_and_special():
    # Edge: unit with underscore and special chars
    code_locations = [
        (("counter", "x", "sec_ond!@#"), {"file": "y.py"})
    ]
    result = decode_result(_encode_locations(0, code_locations)) # 16.0μs -> 16.2μs (0.798% slower)

def test_empty_metric_name_and_unit():
    # Edge: empty metric name and unit
    code_locations = [
        (("counter", "", ""), {"file": "z.py"})
    ]
    result = decode_result(_encode_locations(0, code_locations)) # 15.3μs -> 14.2μs (7.35% faster)



def test_location_dict_is_shared():
    # Edge: location dict is shared between input and output (mutated)
    loc_dict = {"file": "main.py"}
    code_locations = [(("counter", "requests", "second"), loc_dict)]
    _encode_locations(0, code_locations) # 22.1μs -> 21.1μs (5.05% faster)

def test_unit_with_only_special_chars():
    # Edge: unit is only special chars, should be sanitized to empty string
    code_locations = [
        (("counter", "requests", "@#$%^&*()"), {"file": "main.py"})
    ]
    result = decode_result(_encode_locations(0, code_locations)) # 17.7μs -> 17.3μs (2.12% faster)

def test_metric_name_with_multiple_special():
    # Edge: metric name with multiple consecutive special chars
    code_locations = [
        (("counter", "re@@@quests!!!", "second"), {"file": "main.py"})
    ]
    result = decode_result(_encode_locations(0, code_locations)) # 17.4μs -> 16.9μs (2.95% faster)

def test_location_dict_is_copied_for_each():
    # Edge: multiple locations with same dict object, both mutated
    loc_dict = {"file": "main.py"}
    code_locations = [
        (("counter", "requests", "second"), loc_dict),
        (("counter", "requests", "second"), loc_dict),
    ]
    _encode_locations(0, code_locations) # 19.0μs -> 18.2μs (4.79% faster)

# 3. Large Scale Test Cases




def test_large_scale_performance():
    # Large Scale: ensure function does not crash or hang with 999 locations
    N = 999
    code_locations = [
        (("counter", f"metric{i}", "unit"), {"file": f"file{i}.py"})
        for i in range(N)
    ]
    # Just ensure it returns bytes and decodes to dict
    codeflash_output = _encode_locations(0, code_locations); result = codeflash_output # 2.46ms -> 2.17ms (13.2% faster)
    decoded = decode_result(result)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
import json
import re
from functools import partial
from typing import Any, Dict, Iterable, List, Tuple

# imports
import pytest  # used for our unit tests
from sentry_sdk.metrics import _encode_locations

_sanitize_unit = partial(re.compile(r"[^a-zA-Z0-9_]+").sub, "")
_sanitize_metric_key = partial(re.compile(r"[^a-zA-Z0-9_\-.]+").sub, "_")
from sentry_sdk.metrics import _encode_locations

# unit tests

# ---------- BASIC TEST CASES ----------

def test_basic_single_location():
    # Test with a single location and simple metric key
    timestamp = 1234567890
    code_locations = [
        (("c", "my_metric", "s"), {"file": "main.py", "line": 42})
    ]
    codeflash_output = _encode_locations(timestamp, code_locations); result = codeflash_output # 22.2μs -> 21.3μs (4.19% faster)
    decoded = json.loads(result.decode("utf-8"))
    # Check location dict
    loc = decoded["mapping"]["c:my_metric@s"][0]

def test_basic_multiple_locations_same_key():
    # Multiple locations for the same metric key
    timestamp = 1
    code_locations = [
        (("g", "foo", "ms"), {"file": "a.py", "line": 1}),
        (("g", "foo", "ms"), {"file": "b.py", "line": 2}),
    ]
    codeflash_output = _encode_locations(timestamp, code_locations); result = codeflash_output # 20.6μs -> 19.5μs (6.02% faster)
    decoded = json.loads(result.decode("utf-8"))
    locs = decoded["mapping"]["g:foo@ms"]

def test_basic_multiple_keys():
    # Multiple metric keys and locations
    timestamp = 2
    code_locations = [
        (("c", "metricA", "s"), {"file": "x.py", "line": 10}),
        (("g", "metricB", "ms"), {"file": "y.py", "line": 20}),
    ]
    codeflash_output = _encode_locations(timestamp, code_locations); result = codeflash_output # 20.7μs -> 20.3μs (2.05% faster)
    decoded = json.loads(result.decode("utf-8"))

def test_basic_empty_locations():
    # Empty code_locations should result in empty mapping
    timestamp = 3
    code_locations = []
    codeflash_output = _encode_locations(timestamp, code_locations); result = codeflash_output # 10.9μs -> 11.0μs (0.455% slower)
    decoded = json.loads(result.decode("utf-8"))

# ---------- EDGE TEST CASES ----------


def test_edge_empty_strings_in_metric_key_and_unit():
    # Empty name and unit strings should be handled correctly
    timestamp = 100
    code_locations = [
        (("c", "", ""), {"file": "empty.py", "line": 0}),
    ]
    codeflash_output = _encode_locations(timestamp, code_locations); result = codeflash_output # 21.3μs -> 18.8μs (13.7% faster)
    decoded = json.loads(result.decode("utf-8"))
    # name: "" -> "" (no change, but sanitized to "")
    # unit: "" -> "" (no change)
    expected_key = "c:@"
    loc = decoded["mapping"][expected_key][0]

def test_edge_non_ascii_characters():
    # Non-ASCII characters in metric key and unit should be sanitized
    timestamp = 555
    code_locations = [
        (("g", "métrïc💡", "µs"), {"file": "unicode.py", "line": 8}),
    ]
    codeflash_output = _encode_locations(timestamp, code_locations); result = codeflash_output # 20.8μs -> 20.4μs (1.61% faster)
    decoded = json.loads(result.decode("utf-8"))
    # name: "métrïc💡" -> "m_tr_c_"
    # unit: "µs" -> "s"
    expected_key = "g:m_tr_c_@s"
    loc = decoded["mapping"][expected_key][0]

def test_edge_location_dict_is_mutated():
    # The original location dict should be mutated with "type": "location"
    timestamp = 999
    loc_dict = {"file": "mutate.py", "line": 123}
    code_locations = [
        (("c", "mutate", "s"), loc_dict),
    ]
    _encode_locations(timestamp, code_locations) # 16.4μs -> 16.1μs (2.18% faster)

def test_edge_duplicate_metric_keys_with_different_locations():
    # Same metric key, different locations, should all be present
    timestamp = 101
    code_locations = [
        (("c", "dup", "s"), {"file": "a.py", "line": 1}),
        (("c", "dup", "s"), {"file": "b.py", "line": 2}),
        (("c", "dup", "s"), {"file": "c.py", "line": 3}),
    ]
    codeflash_output = _encode_locations(timestamp, code_locations); result = codeflash_output # 22.5μs -> 20.7μs (8.82% faster)
    decoded = json.loads(result.decode("utf-8"))
    key = "c:dup@s"
    files = [loc["file"] for loc in decoded["mapping"][key]]

def test_edge_location_dict_with_extra_fields():
    # Location dict with extra fields should be preserved
    timestamp = 202
    code_locations = [
        (("c", "extra", "s"), {"file": "extra.py", "line": 5, "extra": "value"}),
    ]
    codeflash_output = _encode_locations(timestamp, code_locations); result = codeflash_output # 16.6μs -> 16.2μs (2.69% faster)
    decoded = json.loads(result.decode("utf-8"))
    key = "c:extra@s"
    loc = decoded["mapping"][key][0]

def test_edge_location_dict_is_shared_reference():
    # If same dict is used for multiple metric keys, it should be mutated for both
    timestamp = 303
    shared_loc = {"file": "shared.py", "line": 7}
    code_locations = [
        (("c", "key1", "s"), shared_loc),
        (("g", "key2", "ms"), shared_loc),
    ]
    _encode_locations(timestamp, code_locations) # 20.0μs -> 19.5μs (2.58% faster)

def test_edge_location_dict_is_not_copied():
    # Ensure the function does not copy the location dict
    timestamp = 404
    loc = {"file": "nocopy.py", "line": 11}
    code_locations = [
        (("c", "nocopy", "s"), loc),
    ]
    _encode_locations(timestamp, code_locations) # 16.3μs -> 15.7μs (3.89% faster)

def test_edge_metric_key_with_dash_and_dot():
    # Metric key with dash and dot should be sanitized correctly
    timestamp = 505
    code_locations = [
        (("c", "key-name.with.dot", "s"), {"file": "dashdot.py", "line": 6}),
    ]
    codeflash_output = _encode_locations(timestamp, code_locations); result = codeflash_output # 16.3μs -> 15.9μs (2.59% faster)
    decoded = json.loads(result.decode("utf-8"))
    # name: "key-name.with.dot" -> "key-name.with.dot"
    key = "c:key-name.with.dot@s"

# ---------- LARGE SCALE TEST CASES ----------

def test_large_scale_many_locations():
    # Test with a large number of locations (up to 1000)
    timestamp = 888
    code_locations = [
        (("c", f"metric{i}", "s"), {"file": f"file{i}.py", "line": i})
        for i in range(1000)
    ]
    codeflash_output = _encode_locations(timestamp, code_locations); result = codeflash_output # 2.51ms -> 2.29ms (9.73% faster)
    decoded = json.loads(result.decode("utf-8"))
    # All keys should be present
    for i in range(1000):
        key = f"c:metric{i}@s"
        loc = decoded["mapping"][key][0]

def test_large_scale_many_locations_same_key():
    # Test with 1000 locations, all with the same metric key
    timestamp = 777
    code_locations = [
        (("g", "same", "ms"), {"file": f"file{i}.py", "line": i})
        for i in range(1000)
    ]
    codeflash_output = _encode_locations(timestamp, code_locations); result = codeflash_output # 2.11ms -> 1.85ms (13.7% faster)
    decoded = json.loads(result.decode("utf-8"))
    key = "g:same@ms"
    locs = decoded["mapping"][key]
    for i in range(1000):
        pass

def test_large_scale_varied_metric_keys_and_units():
    # Test with 1000 locations, varied metric keys and units
    timestamp = 666
    code_locations = [
        ((mt, f"name{i}", unit), {"file": f"file{i}.py", "line": i})
        for i, (mt, unit) in enumerate(
            zip(
                ["c", "g", "d", "e"] * 250,
                ["s", "ms", "kg", "m"] * 250
            )
        )
    ]
    codeflash_output = _encode_locations(timestamp, code_locations); result = codeflash_output # 2.51ms -> 2.27ms (10.2% faster)
    decoded = json.loads(result.decode("utf-8"))
    # Check a few keys
    for i in [0, 499, 999]:
        mt = ["c", "g", "d", "e"][i % 4]
        unit = ["s", "ms", "kg", "m"][i % 4]
        key = f"{mt}:name{i}@{unit}"
        loc = decoded["mapping"][key][0]

To edit these changes git checkout codeflash/optimize-_encode_locations-mg9839pe and push.