refactor: replace algorithm dict lookup with match/case pattern
This commit is contained in:
@@ -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
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user