⚡️ Speed up function update_openapi by 41%

[codeflash-ai]

📄 41% (0.41x) speedup for update_openapi in src/titiler/core/titiler/core/utils.py

⏱️ Runtime : 77.1 microseconds → 54.5 microseconds (best of 9 runs)

📝 Explanation and details

The optimized code achieves a 41% speedup by replacing the next() generator expression with an explicit for loop.

Key optimization:

• Eliminated generator overhead: The original code uses next(route for route in app.router.routes if route.path == app.openapi_url) which creates a generator object and involves Python's generator machinery. The optimized version uses a simple for loop with early termination via break, avoiding the generator creation and iteration overhead.

Why this works:

• Generator expressions in Python have overhead for creation and the next() function call
• A direct for loop with break is more efficient for finding the first matching item
• The loop avoids the intermediate generator object allocation
• Early termination with break ensures we don't iterate through remaining routes unnecessarily

Performance characteristics:

• Best case: When the OpenAPI route is found early in the routes list (63-68% faster in large-scale tests)
• Consistent improvement: Shows 30-50% speedup across all test scenarios
• Scales well: Larger route collections see greater benefits (up to 68% improvement with 1000 routes)

The optimization maintains identical behavior including raising StopIteration when no matching route is found, making it a drop-in replacement with pure performance benefits. This is particularly valuable for applications with many routes where the OpenAPI endpoint setup happens during application initialization.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import pytest  # used for our unit tests
from fastapi import FastAPI
from fastapi.testclient import TestClient
from starlette.requests import Request
from starlette.responses import Response
from titiler.core.utils import update_openapi

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

def test_openapi_content_type_basic():
    """
    Basic: Ensure that after patching, /openapi.json returns the correct content-type header.
    """
    # Create a simple FastAPI app
    app = FastAPI()
    # Patch the OpenAPI route
    update_openapi(app) # 3.26μs -> 2.26μs (44.4% faster)
    # Use TestClient to call the openapi route
    client = TestClient(app)
    resp = client.get(app.openapi_url)

def test_openapi_content_type_with_custom_openapi_url():
    """
    Basic: Ensure patching works if the app uses a custom openapi_url.
    """
    app = FastAPI(openapi_url="/docs/openapi.json")
    update_openapi(app) # 3.19μs -> 2.27μs (40.6% faster)
    client = TestClient(app)
    resp = client.get(app.openapi_url)

def test_update_openapi_returns_app():
    """
    Basic: Ensure that update_openapi returns the same app object.
    """
    app = FastAPI()
    codeflash_output = update_openapi(app); patched_app = codeflash_output # 3.34μs -> 2.19μs (52.8% faster)

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

def test_openapi_content_type_not_overwritten_other_routes():
    """
    Edge: Ensure that other routes' content-type headers are not affected.
    """
    app = FastAPI()
    @app.get("/hello")
    def hello():
        return {"msg": "hi"}
    update_openapi(app) # 2.85μs -> 2.07μs (37.5% faster)
    client = TestClient(app)
    resp = client.get("/hello")

def test_openapi_route_missing_raises():
    """
    Edge: If the app has no openapi route, update_openapi should raise StopIteration.
    """
    app = FastAPI(openapi_url=None)
    # Remove openapi route by setting openapi_url to None
    with pytest.raises(StopIteration):
        update_openapi(app) # 1.98μs -> 1.45μs (36.5% faster)

def test_openapi_route_with_additional_headers():
    """
    Edge: Ensure that additional headers set by FastAPI are preserved except content-type.
    """
    app = FastAPI()
    update_openapi(app) # 2.92μs -> 2.11μs (38.1% faster)
    client = TestClient(app)
    resp = client.get(app.openapi_url)

def test_openapi_content_type_multiple_patchings():
    """
    Edge: Ensure that patching multiple times does not break the route.
    """
    app = FastAPI()
    update_openapi(app) # 3.14μs -> 2.21μs (42.2% faster)
    update_openapi(app) # 1.44μs -> 1.20μs (19.5% faster)
    client = TestClient(app)
    resp = client.get(app.openapi_url)

def test_openapi_content_type_case_sensitivity():
    """
    Edge: Ensure that header name is always lowercased as per HTTP spec.
    """
    app = FastAPI()
    update_openapi(app) # 3.08μs -> 2.13μs (44.3% faster)
    client = TestClient(app)
    resp = client.get(app.openapi_url)

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

def test_openapi_content_type_with_many_routes():
    """
    Large Scale: Ensure update_openapi works when app has many routes.
    """
    app = FastAPI()
    # Add 500 routes
    for i in range(500):
        app.get(f"/route{i}")(lambda: {"i": i})
    update_openapi(app) # 4.73μs -> 2.90μs (63.0% faster)
    client = TestClient(app)
    resp = client.get(app.openapi_url)
    # Check that all routes still work
    for i in range(0, 500, 100):
        r = client.get(f"/route{i}")

def test_openapi_content_type_with_large_openapi_schema():
    """
    Large Scale: Ensure update_openapi works when OpenAPI schema is large.
    """
    app = FastAPI()
    # Add 1000 endpoints with unique tags to bloat the OpenAPI schema
    for i in range(1000):
        app.get(f"/item{i}", tags=[f"tag{i}"])(lambda: {"item": i})
    update_openapi(app) # 4.89μs -> 2.91μs (68.1% faster)
    client = TestClient(app)
    resp = client.get(app.openapi_url)
    # OpenAPI schema should contain all tags
    tags = [t["name"] for t in resp.json().get("tags", [])]

def test_openapi_content_type_performance():
    """
    Large Scale: Ensure update_openapi does not significantly slow down openapi route.
    """
    import time
    app = FastAPI()
    # Add 500 endpoints
    for i in range(500):
        app.get(f"/endpoint{i}")(lambda: {"val": i})
    update_openapi(app) # 4.63μs -> 2.99μs (54.8% faster)
    client = TestClient(app)
    # Time the openapi route
    start = time.time()
    resp = client.get(app.openapi_url)
    duration = time.time() - start
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

import pytest
from fastapi import FastAPI
from fastapi.testclient import TestClient
# function to test
from starlette.requests import Request
from starlette.responses import Response
from starlette.routing import Route, request_response
from titiler.core.utils import update_openapi

# unit tests

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

def test_openapi_content_type_replacement_custom_openapi_url():
    """Basic test: App with custom openapi_url."""
    app = FastAPI(openapi_url="/docs/openapi.json")
    client = TestClient(update_openapi(app)) # 3.21μs -> 2.26μs (41.9% faster)
    resp = client.get("/docs/openapi.json")

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

def test_openapi_content_type_already_set():
    """Edge: If content-type is already correct, ensure it's not duplicated or malformed."""
    # Patch the app, then patch again (should be idempotent)
    app = FastAPI()
    update_openapi(app) # 2.71μs -> 2.07μs (30.7% faster)
    update_openapi(app) # 1.70μs -> 1.18μs (43.9% faster)
    client = TestClient(app)
    resp = client.get(app.openapi_url)

def test_openapi_route_not_found():
    """Edge: If openapi_url is set to a non-existent route, function should raise StopIteration."""
    app = FastAPI(openapi_url="/notfound.json")
    # Remove the openapi route to simulate missing route
    app.router.routes = [r for r in app.router.routes if getattr(r, "path", None) != "/notfound.json"]
    with pytest.raises(StopIteration):
        update_openapi(app) # 2.13μs -> 1.68μs (26.7% faster)

def test_openapi_route_with_additional_headers():
    """Edge: If the OpenAPI route adds its own headers, ensure our header overwrites content-type only."""
    app = FastAPI()

    # Patch the OpenAPI endpoint to add a custom header before patching
    openapi_route = next(route for route in app.router.routes if route.path == app.openapi_url)
    old_endpoint = openapi_route.endpoint

    async def custom_header_endpoint(request: Request):
        response = await old_endpoint(request)
        response.headers["X-Custom"] = "foobar"
        return response

    openapi_route.app = request_response(custom_header_endpoint)

    update_openapi(app) # 2.08μs -> 1.81μs (14.8% faster)
    client = TestClient(app)
    resp = client.get(app.openapi_url)

def test_openapi_route_method_is_get_only():
    """Edge: Ensure the OpenAPI route only responds to GET, and is not broken by patch."""
    app = FastAPI()
    update_openapi(app) # 3.15μs -> 2.24μs (40.5% faster)
    client = TestClient(app)
    # GET should work
    resp = client.get(app.openapi_url)
    # POST should not be allowed
    resp_post = client.post(app.openapi_url)

def test_openapi_content_type_case_insensitive_header():
    """Edge: HTTP headers are case-insensitive, ensure content-type is set correctly."""
    app = FastAPI()
    update_openapi(app) # 3.07μs -> 2.08μs (47.7% faster)
    client = TestClient(app)
    resp = client.get(app.openapi_url)

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

def test_openapi_with_many_routes():
    """Large scale: App with hundreds of routes, OpenAPI patch still works."""
    app = FastAPI()
    # Add 500 routes
    for i in range(500):
        app.get(f"/route{i}")(lambda i=i: {"route": i})
    update_openapi(app) # 3.70μs -> 2.96μs (24.9% faster)
    client = TestClient(app)
    resp = client.get(app.openapi_url)
    data = resp.json()
    # Ensure a sample of our routes are present in OpenAPI
    for i in (0, 100, 250, 499):
        pass

def test_openapi_performance_large_schema(benchmark):
    """Large scale: Benchmark OpenAPI route with many endpoints."""
    app = FastAPI()
    # Add 1000 routes
    for i in range(1000):
        app.get(f"/item/{i}")(lambda i=i: {"item": i})
    update_openapi(app) # 4.11μs -> 3.26μs (25.9% faster)
    client = TestClient(app)
    # Benchmark the OpenAPI route
    def get_openapi():
        resp = client.get(app.openapi_url)
        data = resp.json()
    benchmark(get_openapi)

def test_openapi_patch_does_not_affect_other_routes():
    """Large scale: Ensure patching OpenAPI does not affect normal endpoints."""
    app = FastAPI()
    @app.get("/foo")
    def foo():
        return {"foo": "bar"}
    update_openapi(app) # 3.73μs -> 2.60μs (43.6% faster)
    client = TestClient(app)
    # /foo should still work and have normal content-type
    resp = client.get("/foo")

# ---- Misc/Robustness ----

def test_update_openapi_returns_same_app_object():
    """Misc: Ensure the returned app is the same object as input."""
    app = FastAPI()
    codeflash_output = update_openapi(app); app2 = codeflash_output # 3.52μs -> 2.42μs (45.6% faster)

def test_update_openapi_multiple_times_is_idempotent():
    """Misc: Calling update_openapi multiple times is safe and does not break OpenAPI."""
    app = FastAPI()
    update_openapi(app) # 2.87μs -> 2.01μs (42.7% faster)
    update_openapi(app) # 1.65μs -> 1.20μs (37.5% faster)
    client = TestClient(app)
    resp = client.get(app.openapi_url)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

To edit these changes git checkout codeflash/optimize-update_openapi-mi8d8bbm and push.