⚡️ Speed up function get_role_based_routes by 24%

[codeflash-ai]

📄 24% (0.24x) speedup for get_role_based_routes in litellm/proxy/auth/auth_checks.py

⏱️ Runtime : 239 microseconds → 192 microseconds (best of 60 runs)

📝 Explanation and details

The optimization achieves a 24% speedup by replacing a manual for-loop with Python's built-in next() function and a generator expression, along with removing unnecessary type casting.

Key optimizations applied:

1. Eliminated unnecessary cast() operation: Removed the cast(Optional[List[RoleBasedPermissions]], ...) which added overhead without providing runtime benefit.

2. Replaced for-loop with next() + generator: Changed from manually iterating through role_based_permissions to using next((rbp for rbp in role_based_permissions if rbp.role == rbac_role), None). This leverages Python's optimized C-level implementation and enables short-circuit evaluation - stopping immediately when the first match is found.

3. Simplified truthiness check: Changed if role_based_permissions is None: to if not role_based_permissions: which handles both None and empty list cases more efficiently.

Why this is faster:

• The next() function with generator expression is implemented in C and optimized for early termination
• Removes the overhead of explicit loop iteration and condition checking in Python bytecode
• The line profiler shows the optimization is particularly effective when the target role appears early in the list or when no match is found

Impact on workloads:
Based on the function reference, get_role_based_routes() is called from can_rbac_role_call_route() in the JWT authentication flow. This means the optimization affects every authenticated request that uses role-based access control. The test results show especially significant gains (up to 172% faster) for edge cases like empty permissions lists, which are common in misconfigured or initialization scenarios.

Test case performance patterns:

• Best gains (140-170% faster): Empty/None permission lists - common during system startup or misconfiguration
• Good gains (15-40% faster): Role lookup scenarios where the target role is found early or not found at all
• Consistent gains (10-25% faster): Normal operation with populated permission lists

This optimization is particularly valuable since authentication happens on the critical path of every API request.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

from typing import List, Optional

# imports
import pytest
from litellm.proxy.auth.auth_checks import get_role_based_routes

# --- Function and supporting classes/types to test ---

# Simulate RBAC_ROLES and RoleBasedPermissions for testing purposes
class RBAC_ROLES:
    ADMIN = "admin"
    USER = "user"
    GUEST = "guest"
    SUPERUSER = "superuser"
    EMPTY = ""
    NONE = None
    CUSTOM = "custom"

class RoleBasedPermissions:
    def __init__(self, role: str, models: Optional[List[str]] = None, routes: Optional[List[str]] = None):
        self.role = role
        self.models = models if models is not None else []
        self.routes = routes if routes is not None else []
from litellm.proxy.auth.auth_checks import get_role_based_routes

# --- Unit tests ---

# 1. Basic Test Cases
def test_basic_admin_routes():
    # Admin role has access to /admin and /dashboard
    general_settings = {
        "role_permissions": [
            RoleBasedPermissions(role=RBAC_ROLES.ADMIN, routes=["/admin", "/dashboard"])
        ]
    }
    codeflash_output = get_role_based_routes(RBAC_ROLES.ADMIN, general_settings) # 2.95μs -> 2.48μs (19.2% faster)

def test_basic_user_routes():
    # User role has access to /home and /profile
    general_settings = {
        "role_permissions": [
            RoleBasedPermissions(role=RBAC_ROLES.USER, routes=["/home", "/profile"])
        ]
    }
    codeflash_output = get_role_based_routes(RBAC_ROLES.USER, general_settings) # 2.72μs -> 2.19μs (24.3% faster)

def test_basic_guest_routes():
    # Guest role has access to just /home
    general_settings = {
        "role_permissions": [
            RoleBasedPermissions(role=RBAC_ROLES.GUEST, routes=["/home"])
        ]
    }
    codeflash_output = get_role_based_routes(RBAC_ROLES.GUEST, general_settings) # 2.63μs -> 2.06μs (27.5% faster)

def test_basic_role_not_in_permissions():
    # Role not present in permissions returns None
    general_settings = {
        "role_permissions": [
            RoleBasedPermissions(role=RBAC_ROLES.ADMIN, routes=["/admin"]),
            RoleBasedPermissions(role=RBAC_ROLES.USER, routes=["/user"])
        ]
    }
    codeflash_output = get_role_based_routes("unknown_role", general_settings) # 2.45μs -> 1.74μs (40.8% faster)

def test_basic_multiple_roles():
    # Multiple roles, check that correct one is selected
    general_settings = {
        "role_permissions": [
            RoleBasedPermissions(role=RBAC_ROLES.ADMIN, routes=["/admin"]),
            RoleBasedPermissions(role=RBAC_ROLES.USER, routes=["/user"]),
            RoleBasedPermissions(role=RBAC_ROLES.GUEST, routes=["/guest"])
        ]
    }
    codeflash_output = get_role_based_routes(RBAC_ROLES.USER, general_settings) # 2.89μs -> 2.28μs (26.7% faster)

# 2. Edge Test Cases

def test_edge_empty_role_permissions_list():
    # Empty role_permissions list returns None for any role
    general_settings = {"role_permissions": []}
    codeflash_output = get_role_based_routes(RBAC_ROLES.ADMIN, general_settings) # 2.32μs -> 894ns (160% faster)

def test_edge_role_permissions_none():
    # role_permissions is None returns None for any role
    general_settings = {"role_permissions": None}
    codeflash_output = get_role_based_routes(RBAC_ROLES.ADMIN, general_settings) # 2.09μs -> 816ns (156% faster)

def test_edge_missing_role_permissions_key():
    # role_permissions key missing should default to empty list, returns None
    general_settings = {}
    codeflash_output = get_role_based_routes(RBAC_ROLES.ADMIN, general_settings) # 2.31μs -> 875ns (164% faster)

def test_edge_role_with_no_routes():
    # Role exists but routes is empty
    general_settings = {
        "role_permissions": [
            RoleBasedPermissions(role=RBAC_ROLES.ADMIN, routes=[])
        ]
    }
    codeflash_output = get_role_based_routes(RBAC_ROLES.ADMIN, general_settings) # 2.57μs -> 2.35μs (9.27% faster)

def test_edge_role_with_none_routes():
    # Role exists but routes is None
    general_settings = {
        "role_permissions": [
            RoleBasedPermissions(role=RBAC_ROLES.ADMIN, routes=None)
        ]
    }
    codeflash_output = get_role_based_routes(RBAC_ROLES.ADMIN, general_settings) # 2.56μs -> 2.23μs (14.7% faster)

def test_edge_role_is_empty_string():
    # Role is empty string
    general_settings = {
        "role_permissions": [
            RoleBasedPermissions(role="", routes=["/empty"])
        ]
    }
    codeflash_output = get_role_based_routes("", general_settings) # 2.46μs -> 2.05μs (20.0% faster)

def test_edge_role_is_none():
    # Role is None
    general_settings = {
        "role_permissions": [
            RoleBasedPermissions(role=None, routes=["/none"])
        ]
    }
    codeflash_output = get_role_based_routes(None, general_settings) # 2.62μs -> 2.12μs (23.4% faster)

def test_edge_duplicate_roles():
    # Duplicate roles, should return first match
    general_settings = {
        "role_permissions": [
            RoleBasedPermissions(role=RBAC_ROLES.USER, routes=["/first"]),
            RoleBasedPermissions(role=RBAC_ROLES.USER, routes=["/second"])
        ]
    }
    codeflash_output = get_role_based_routes(RBAC_ROLES.USER, general_settings) # 2.45μs -> 2.13μs (14.9% faster)

def test_edge_routes_are_empty_strings():
    # Routes are empty strings
    general_settings = {
        "role_permissions": [
            RoleBasedPermissions(role=RBAC_ROLES.ADMIN, routes=[""])
        ]
    }
    codeflash_output = get_role_based_routes(RBAC_ROLES.ADMIN, general_settings) # 2.39μs -> 2.01μs (18.5% faster)

def test_edge_routes_are_none_in_list():
    # Routes list contains None values
    general_settings = {
        "role_permissions": [
            RoleBasedPermissions(role=RBAC_ROLES.ADMIN, routes=[None, "/admin"])
        ]
    }
    codeflash_output = get_role_based_routes(RBAC_ROLES.ADMIN, general_settings) # 2.35μs -> 1.97μs (19.2% faster)



def test_large_scale_many_roles():
    # 1000 roles, select one in the middle
    num_roles = 1000
    target_index = 500
    roles = [f"role_{i}" for i in range(num_roles)]
    routes = [f"/route_{i}" for i in range(num_roles)]
    general_settings = {
        "role_permissions": [
            RoleBasedPermissions(role=role, routes=[route])
            for role, route in zip(roles, routes)
        ]
    }
    codeflash_output = get_role_based_routes(roles[target_index], general_settings) # 21.1μs -> 18.4μs (14.8% faster)

def test_large_scale_many_routes_for_one_role():
    # One role with 1000 routes
    many_routes = [f"/route_{i}" for i in range(1000)]
    general_settings = {
        "role_permissions": [
            RoleBasedPermissions(role=RBAC_ROLES.SUPERUSER, routes=many_routes)
        ]
    }
    codeflash_output = get_role_based_routes(RBAC_ROLES.SUPERUSER, general_settings) # 3.06μs -> 2.63μs (16.2% faster)

def test_large_scale_all_roles_none_except_one():
    # 999 roles with None, one with routes
    num_roles = 1000
    general_settings = {
        "role_permissions": [
            RoleBasedPermissions(role=f"role_{i}", routes=None) for i in range(num_roles - 1)
        ] + [RoleBasedPermissions(role="target", routes=["/target"])]
    }
    codeflash_output = get_role_based_routes("target", general_settings) # 33.7μs -> 29.1μs (15.8% faster)

def test_large_scale_role_at_end_of_list():
    # Role is last in the list
    num_roles = 1000
    general_settings = {
        "role_permissions": [
            RoleBasedPermissions(role=f"role_{i}", routes=[f"/route_{i}"]) for i in range(num_roles - 1)
        ] + [RoleBasedPermissions(role="last_role", routes=["/last"])]
    }
    codeflash_output = get_role_based_routes("last_role", general_settings) # 34.1μs -> 29.0μs (17.6% faster)

def test_large_scale_role_permissions_is_empty_list():
    # role_permissions is empty even with large data
    general_settings = {"role_permissions": []}
    codeflash_output = get_role_based_routes("any_role", general_settings) # 2.61μs -> 1.08μs (142% faster)

def test_large_scale_role_permissions_is_none():
    # role_permissions is None even with large data
    general_settings = {"role_permissions": None}
    codeflash_output = get_role_based_routes("any_role", general_settings) # 2.19μs -> 907ns (141% faster)

from typing import List, Optional

# imports
import pytest  # used for our unit tests
from litellm.proxy.auth.auth_checks import get_role_based_routes

# --- Function and dependencies to test ---

# Simulate RBAC_ROLES and RoleBasedPermissions as per the original source
class RBAC_ROLES:
    ADMIN = "admin"
    USER = "user"
    GUEST = "guest"
    SUPERUSER = "superuser"
    DEVELOPER = "developer"
    # Add more roles as needed

class RoleBasedPermissions:
    def __init__(self, role: str, models: Optional[List[str]] = None, routes: Optional[List[str]] = None):
        self.role = role
        self.models = models or []
        self.routes = routes or []
from litellm.proxy.auth.auth_checks import get_role_based_routes

# --- Unit Tests ---

# BASIC TEST CASES

def test_basic_admin_role_with_routes():
    # Test admin role has correct routes
    general_settings = {
        "role_permissions": [
            RoleBasedPermissions(role=RBAC_ROLES.ADMIN, routes=["/dashboard", "/admin"]),
            RoleBasedPermissions(role=RBAC_ROLES.USER, routes=["/home"]),
        ]
    }
    codeflash_output = get_role_based_routes(RBAC_ROLES.ADMIN, general_settings); result = codeflash_output # 3.59μs -> 2.79μs (28.6% faster)

def test_basic_user_role_with_routes():
    # Test user role has correct routes
    general_settings = {
        "role_permissions": [
            RoleBasedPermissions(role=RBAC_ROLES.ADMIN, routes=["/dashboard", "/admin"]),
            RoleBasedPermissions(role=RBAC_ROLES.USER, routes=["/home"]),
        ]
    }
    codeflash_output = get_role_based_routes(RBAC_ROLES.USER, general_settings); result = codeflash_output # 2.91μs -> 2.59μs (12.6% faster)

def test_basic_guest_role_not_present():
    # Test role not present returns None
    general_settings = {
        "role_permissions": [
            RoleBasedPermissions(role=RBAC_ROLES.ADMIN, routes=["/dashboard", "/admin"]),
            RoleBasedPermissions(role=RBAC_ROLES.USER, routes=["/home"]),
        ]
    }
    codeflash_output = get_role_based_routes(RBAC_ROLES.GUEST, general_settings); result = codeflash_output # 2.73μs -> 1.99μs (37.6% faster)

def test_basic_empty_role_permissions():
    # Test with empty role_permissions list
    general_settings = {"role_permissions": []}
    codeflash_output = get_role_based_routes(RBAC_ROLES.ADMIN, general_settings); result = codeflash_output # 2.36μs -> 866ns (172% faster)

def test_basic_none_role_permissions():
    # Test with None role_permissions
    general_settings = {"role_permissions": None}
    codeflash_output = get_role_based_routes(RBAC_ROLES.ADMIN, general_settings); result = codeflash_output # 2.40μs -> 916ns (162% faster)

def test_basic_missing_role_permissions_key():
    # Test with missing role_permissions key
    general_settings = {}
    codeflash_output = get_role_based_routes(RBAC_ROLES.ADMIN, general_settings); result = codeflash_output # 2.48μs -> 927ns (167% faster)

def test_basic_role_with_no_routes():
    # Test role present but routes is empty
    general_settings = {
        "role_permissions": [
            RoleBasedPermissions(role=RBAC_ROLES.ADMIN, routes=[]),
        ]
    }
    codeflash_output = get_role_based_routes(RBAC_ROLES.ADMIN, general_settings); result = codeflash_output # 2.76μs -> 2.41μs (14.7% faster)

# EDGE TEST CASES

def test_edge_multiple_roles_same_name():
    # Test multiple entries for same role, should return first match
    general_settings = {
        "role_permissions": [
            RoleBasedPermissions(role=RBAC_ROLES.USER, routes=["/profile"]),
            RoleBasedPermissions(role=RBAC_ROLES.USER, routes=["/settings"]),
        ]
    }
    codeflash_output = get_role_based_routes(RBAC_ROLES.USER, general_settings); result = codeflash_output # 2.73μs -> 2.29μs (19.1% faster)

def test_edge_role_with_none_routes():
    # Test role with routes explicitly set to None
    class RoleBasedPermissionsWithNone(RoleBasedPermissions):
        def __init__(self, role, models=None, routes=None):
            self.role = role
            self.models = models or []
            self.routes = routes  # Allow explicit None

    general_settings = {
        "role_permissions": [
            RoleBasedPermissionsWithNone(role=RBAC_ROLES.DEVELOPER, routes=None),
        ]
    }
    codeflash_output = get_role_based_routes(RBAC_ROLES.DEVELOPER, general_settings); result = codeflash_output # 2.69μs -> 2.37μs (13.5% faster)

def test_edge_role_with_non_list_routes():
    # Test role with routes not a list (should handle gracefully)
    class RoleBasedPermissionsNonList(RoleBasedPermissions):
        def __init__(self, role, models=None, routes=None):
            self.role = role
            self.models = models or []
            self.routes = routes  # Could be non-list

    general_settings = {
        "role_permissions": [
            RoleBasedPermissionsNonList(role=RBAC_ROLES.ADMIN, routes="/admin"),
        ]
    }
    codeflash_output = get_role_based_routes(RBAC_ROLES.ADMIN, general_settings); result = codeflash_output # 2.75μs -> 2.36μs (16.5% faster)


def test_edge_role_permissions_empty_dict():
    # Test with role_permissions as empty dict
    general_settings = {
        "role_permissions": {}
    }
    codeflash_output = get_role_based_routes(RBAC_ROLES.ADMIN, general_settings); result = codeflash_output # 3.21μs -> 1.25μs (157% faster)

def test_edge_role_is_none():
    # Test with rbac_role as None
    general_settings = {
        "role_permissions": [
            RoleBasedPermissions(role=None, routes=["/none"]),
        ]
    }
    codeflash_output = get_role_based_routes(None, general_settings); result = codeflash_output # 2.98μs -> 2.70μs (10.3% faster)

def test_edge_case_sensitive_role_names():
    # Test role name case sensitivity
    general_settings = {
        "role_permissions": [
            RoleBasedPermissions(role="Admin", routes=["/dashboard"]),
        ]
    }
    codeflash_output = get_role_based_routes(RBAC_ROLES.ADMIN, general_settings); result = codeflash_output # 2.58μs -> 1.87μs (38.3% faster)

def test_edge_role_with_extra_keys_in_permission():
    # Test permission objects with extra keys
    class RoleBasedPermissionsExtra(RoleBasedPermissions):
        def __init__(self, role, models=None, routes=None, extra=None):
            super().__init__(role, models, routes)
            self.extra = extra

    general_settings = {
        "role_permissions": [
            RoleBasedPermissionsExtra(role=RBAC_ROLES.ADMIN, routes=["/admin"], extra="extra_value"),
        ]
    }
    codeflash_output = get_role_based_routes(RBAC_ROLES.ADMIN, general_settings); result = codeflash_output # 2.80μs -> 2.48μs (12.9% faster)

# LARGE SCALE TEST CASES

def test_large_scale_many_roles_and_routes():
    # Test with many roles and many routes
    num_roles = 500
    num_routes = 10
    roles = [f"role_{i}" for i in range(num_roles)]
    route_lists = [[f"/route_{j}" for j in range(num_routes)] for _ in range(num_roles)]
    permissions = [RoleBasedPermissions(role=roles[i], routes=route_lists[i]) for i in range(num_roles)]
    general_settings = {"role_permissions": permissions}

    # Pick a random role to test
    test_idx = 123
    test_role = roles[test_idx]
    expected_routes = route_lists[test_idx]
    codeflash_output = get_role_based_routes(test_role, general_settings); result = codeflash_output # 8.28μs -> 6.80μs (21.8% faster)

def test_large_scale_all_roles_missing():
    # Test with large number of roles, none matching
    num_roles = 500
    roles = [f"role_{i}" for i in range(num_roles)]
    permissions = [RoleBasedPermissions(role=role, routes=[f"/route_{i}"]) for i, role in enumerate(roles)]
    general_settings = {"role_permissions": permissions}
    codeflash_output = get_role_based_routes("non_existent_role", general_settings); result = codeflash_output # 18.6μs -> 15.4μs (21.0% faster)

def test_large_scale_all_roles_same_routes():
    # Test with large number of roles, all with same routes
    num_roles = 500
    routes = ["/common_route"]
    permissions = [RoleBasedPermissions(role=f"role_{i}", routes=routes) for i in range(num_roles)]
    general_settings = {"role_permissions": permissions}
    test_role = "role_250"
    codeflash_output = get_role_based_routes(test_role, general_settings); result = codeflash_output # 11.5μs -> 9.87μs (16.4% faster)

def test_large_scale_empty_routes_lists():
    # Test with large number of roles, all with empty routes
    num_roles = 500
    permissions = [RoleBasedPermissions(role=f"role_{i}", routes=[]) for i in range(num_roles)]
    general_settings = {"role_permissions": permissions}
    test_role = "role_499"
    codeflash_output = get_role_based_routes(test_role, general_settings); result = codeflash_output # 20.0μs -> 17.4μs (15.0% faster)

def test_large_scale_roles_with_varied_routes():
    # Test with large number of roles, each with different routes
    num_roles = 500
    permissions = []
    for i in range(num_roles):
        routes = [f"/route_{i}_{j}" for j in range(i % 10)]
        permissions.append(RoleBasedPermissions(role=f"role_{i}", routes=routes))
    general_settings = {"role_permissions": permissions}
    test_role = "role_100"
    expected_routes = [f"/route_100_{j}" for j in range(100 % 10)]
    codeflash_output = get_role_based_routes(test_role, general_settings); result = codeflash_output # 6.83μs -> 5.63μs (21.3% faster)
# 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-get_role_based_routes-mhwwp3si and push.