Skip to content

Commit

Permalink
chore(utils): allow duplicate values in registry by making reverse lo…
Browse files Browse the repository at this point in the history
…okup optional (#82114)
  • Loading branch information
cathteng authored Dec 14, 2024
1 parent 3273e16 commit 92e31b8
Show file tree
Hide file tree
Showing 2 changed files with 35 additions and 6 deletions.
21 changes: 15 additions & 6 deletions src/sentry/utils/registry.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,9 +15,15 @@ class NoRegistrationExistsError(ValueError):


class Registry(Generic[T]):
def __init__(self):
"""
A simple generic registry that allows for registering and retrieving items by key. Reverse lookup by value is enabled by default.
If you have duplicate values, you may want to disable reverse lookup.
"""

def __init__(self, enable_reverse_lookup=True):
self.registrations: dict[str, T] = {}
self.reverse_lookup: dict[T, str] = {}
self.enable_reverse_lookup = enable_reverse_lookup

def register(self, key: str):
def inner(item: T) -> T:
Expand All @@ -26,13 +32,14 @@ def inner(item: T) -> T:
f"A registration already exists for {key}: {self.registrations[key]}"
)

if item in self.reverse_lookup:
raise AlreadyRegisteredError(
f"A registration already exists for {item}: {self.reverse_lookup[item]}"
)
if self.enable_reverse_lookup:
if item in self.reverse_lookup:
raise AlreadyRegisteredError(
f"A registration already exists for {item}: {self.reverse_lookup[item]}"
)
self.reverse_lookup[item] = key

self.registrations[key] = item
self.reverse_lookup[item] = key

return item

Expand All @@ -44,6 +51,8 @@ def get(self, key: str) -> T:
return self.registrations[key]

def get_key(self, item: T) -> str:
if not self.enable_reverse_lookup:
raise NotImplementedError("Reverse lookup is not enabled")
if item not in self.reverse_lookup:
raise NoRegistrationExistsError(f"No registration exists for {item}")
return self.reverse_lookup[item]
20 changes: 20 additions & 0 deletions tests/sentry/utils/test_registry.py
Original file line number Diff line number Diff line change
Expand Up @@ -33,3 +33,23 @@ def unregistered_func():

test_registry.register("something else")(unregistered_func)
assert test_registry.get("something else") == unregistered_func

def test_allow_duplicate_values(self):
test_registry = Registry[str](enable_reverse_lookup=False)

@test_registry.register("something")
@test_registry.register("something 2")
def registered_func():
pass

assert test_registry.get("something") == registered_func
assert test_registry.get("something 2") == registered_func

with pytest.raises(NoRegistrationExistsError):
test_registry.get("something else")

with pytest.raises(NotImplementedError):
test_registry.get_key(registered_func)

test_registry.register("something else")(registered_func)
assert test_registry.get("something else") == registered_func

0 comments on commit 92e31b8

Please sign in to comment.