refactor: replace algorithm dict lookup with match/case pattern

This commit is contained in:
2026-02-04 01:07:43 +00:00
parent 3510ea564a
commit 064af30d0f

View File

@@ -26,7 +26,7 @@ class Algorithm(str, Enum):
class BaseAlgorithm(ABC):
"""Base class for rate limiting algorithms."""
__slots__ = ("limit", "window_size", "backend", "burst_size")
__slots__ = ("backend", "burst_size", "limit", "window_size")
def __init__(
self,
@@ -450,17 +450,24 @@ def get_algorithm(
burst_size: int | None = None,
) -> BaseAlgorithm:
"""Factory function to create algorithm instances."""
algorithm_map: dict[Algorithm, type[BaseAlgorithm]] = {
Algorithm.TOKEN_BUCKET: TokenBucketAlgorithm,
Algorithm.SLIDING_WINDOW: SlidingWindowAlgorithm,
Algorithm.FIXED_WINDOW: FixedWindowAlgorithm,
Algorithm.LEAKY_BUCKET: LeakyBucketAlgorithm,
Algorithm.SLIDING_WINDOW_COUNTER: SlidingWindowCounterAlgorithm,
}
algorithm_class = algorithm_map.get(algorithm)
if algorithm_class is None:
msg = f"Unknown algorithm: {algorithm}"
raise ValueError(msg)
return algorithm_class(limit, window_size, backend, burst_size=burst_size)
match algorithm:
case Algorithm.TOKEN_BUCKET:
return TokenBucketAlgorithm(
limit, window_size, backend, burst_size=burst_size
)
case Algorithm.SLIDING_WINDOW:
return SlidingWindowAlgorithm(
limit, window_size, backend, burst_size=burst_size
)
case Algorithm.FIXED_WINDOW:
return FixedWindowAlgorithm(
limit, window_size, backend, burst_size=burst_size
)
case Algorithm.LEAKY_BUCKET:
return LeakyBucketAlgorithm(
limit, window_size, backend, burst_size=burst_size
)
case Algorithm.SLIDING_WINDOW_COUNTER:
return SlidingWindowCounterAlgorithm(
limit, window_size, backend, burst_size=burst_size
)