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) 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. 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. 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(): def test_item_storage_memory_initializes():
item_storage_memory = ItemStorageMemory() item_storage_memory = ItemStorageMemory[MockItemModel]()
assert item_storage_memory._items == {} assert item_storage_memory._items == {}
assert item_storage_memory._id_field == "id" assert item_storage_memory._id_field == "id"
assert item_storage_memory._max_items == 10 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._id_field == "bananas"
assert item_storage_memory._max_items == 20 assert item_storage_memory._max_items == 20
with pytest.raises(ValueError, match=re.escape("max_items must be at least 1")): 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")): 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]): 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(): 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): for i in range(10):
item_storage_memory.set(MockItemModel(id=str(i), value=i)) item_storage_memory.set(MockItemModel(id=str(i), value=i))
assert item_storage_memory._items == { assert item_storage_memory._items == {