chore(item-storage): improve types

Provide type args to the generics.
This commit is contained in:
psychedelicious 2024-02-03 18:35:46 +11:00 committed by Kent Keirsey
parent 88c08bbfc7
commit 9f274c79dc
2 changed files with 6 additions and 6 deletions

View File

@ -10,7 +10,7 @@ from invokeai.app.services.item_storage.item_storage_common import ItemNotFoundE
T = TypeVar("T", bound=BaseModel)
class ItemStorageMemory(ItemStorageABC, Generic[T]):
class ItemStorageMemory(ItemStorageABC[T], Generic[T]):
"""
Provides a simple in-memory storage for items, with a maximum number of items to store.
The storage uses the LRU strategy to evict items from storage when the max has been reached.

View File

@ -18,19 +18,19 @@ def item_storage_memory():
def test_item_storage_memory_initializes():
item_storage_memory = ItemStorageMemory()
item_storage_memory = ItemStorageMemory[MockItemModel]()
assert item_storage_memory._items == {}
assert item_storage_memory._id_field == "id"
assert item_storage_memory._max_items == 10
item_storage_memory = ItemStorageMemory(id_field="bananas", max_items=20)
item_storage_memory = ItemStorageMemory[MockItemModel](id_field="bananas", max_items=20)
assert item_storage_memory._id_field == "bananas"
assert item_storage_memory._max_items == 20
with pytest.raises(ValueError, match=re.escape("max_items must be at least 1")):
item_storage_memory = ItemStorageMemory(max_items=0)
item_storage_memory = ItemStorageMemory[MockItemModel](max_items=0)
with pytest.raises(ValueError, match=re.escape("id_field must not be empty")):
item_storage_memory = ItemStorageMemory(id_field="")
item_storage_memory = ItemStorageMemory[MockItemModel](id_field="")
def test_item_storage_memory_sets(item_storage_memory: ItemStorageMemory[MockItemModel]):
@ -74,7 +74,7 @@ def test_item_storage_memory_deletes(item_storage_memory: ItemStorageMemory[Mock
def test_item_storage_memory_respects_max():
item_storage_memory = ItemStorageMemory(max_items=3)
item_storage_memory = ItemStorageMemory[MockItemModel](max_items=3)
for i in range(10):
item_storage_memory.set(MockItemModel(id=str(i), value=i))
assert item_storage_memory._items == {