diff --git a/src/dexorder/blockstate/blockdata.py b/src/dexorder/blockstate/blockdata.py index dd20741..79cfe3c 100644 --- a/src/dexorder/blockstate/blockdata.py +++ b/src/dexorder/blockstate/blockdata.py @@ -1,7 +1,6 @@ import logging -from collections import defaultdict from enum import Enum -from typing import TypeVar, Generic, Iterable, Union, Any, Iterator +from typing import TypeVar, Generic, Iterable, Union, Any, Iterator, Callable from dexorder import NARG, DELETE, UNLOAD from dexorder.base.fork import current_fork @@ -23,7 +22,6 @@ class DataType(Enum): class BlockData: registry: dict[Any,'BlockData'] = {} # series name and instance - adapters: dict[list['BlockDataAdapter']] = defaultdict(list) def __init__(self, data_type: DataType, series: Any, *, series2str=None, series2key=None, # defaults to key2str and str2key @@ -136,6 +134,17 @@ class BlockDict(Generic[K,V], BlockData): def get(self, item: K, default: V = None) -> V: return self.getitem(item, default) + def modify(self, item: K, func: Callable[[V],V], default: V=NARG) -> V: + try: + result = func(self.getitem(item, default)) + except KeyError: + result = func(NARG) + self.setitem(item, result) + return result + + def add(self, item: K, value: Union[V], default: V = 0) -> Union[V,DELETE]: + return self.modify(item, lambda v: v+value, default=default) + class SeriesCollection: def __init__(self, series_or_datavars: Iterable[Union[Any,BlockData]]):