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