diff --git a/docs/library/btree.rst b/docs/library/btree.rst index 9322d32e6a318..bc07ca5b68406 100644 --- a/docs/library/btree.rst +++ b/docs/library/btree.rst @@ -82,8 +82,10 @@ Functions other parameters are optional and keyword-only, and allow to tweak advanced parameters of the database operation (most users will not need them): - * *flags* - Currently unused. - * *cachesize* - Suggested maximum memory cache size in bytes. For a + * `flags` - 1 to allow duplicate keys (similar to list insertion). + the insertion order of duplicated keys is not specified. + 0 to disable duplicated keys. + * `cachesize` - Suggested maximum memory cache size in bytes. For a board with enough memory using larger values may improve performance. The value is only a recommendation, the module may use more memory if values set too low. @@ -95,6 +97,9 @@ Functions Returns a BTree object, which implements a dictionary protocol (set of methods), and some additional methods described below. + when duplicated keys are allowed (flags=btree.DUPLICATE_KEY), several call + to __setitem__ with the same key will insert instead of replace the item. + The del operation removes all items corresponding to the specified key. Methods ------- @@ -153,3 +158,7 @@ Constants A flag for `keys()`, `values()`, `items()` methods to specify that scanning should be in descending direction of keys. + +.. data:: DUPLICATE_KEY + + A flag for `open` method to allow duplicated keys in the btree. diff --git a/extmod/modbtree.c b/extmod/modbtree.c index 127dd71a395b2..1a8aac3ddd768 100644 --- a/extmod/modbtree.c +++ b/extmod/modbtree.c @@ -46,6 +46,7 @@ typedef struct _mp_obj_btree_t { mp_obj_t end_key; #define FLAG_END_KEY_INCL 1 #define FLAG_DESC 2 + #define FLAG_DUPLICATE_KEY 1 #define FLAG_ITER_TYPE_MASK 0xc0 #define FLAG_ITER_KEYS 0x40 #define FLAG_ITER_VALUES 0x80 @@ -367,6 +368,7 @@ STATIC const mp_rom_map_elem_t mp_module_btree_globals_table[] = { { MP_ROM_QSTR(MP_QSTR_open), MP_ROM_PTR(&mod_btree_open_obj) }, { MP_ROM_QSTR(MP_QSTR_INCL), MP_ROM_INT(FLAG_END_KEY_INCL) }, { MP_ROM_QSTR(MP_QSTR_DESC), MP_ROM_INT(FLAG_DESC) }, + { MP_ROM_QSTR(MP_QSTR_DUPLICATE_KEY), MP_ROM_INT(FLAG_DUPLICATE_KEY) }, }; STATIC MP_DEFINE_CONST_DICT(mp_module_btree_globals, mp_module_btree_globals_table); diff --git a/tests/extmod/btree_dup.py b/tests/extmod/btree_dup.py new file mode 100644 index 0000000000000..f753fb41d3eca --- /dev/null +++ b/tests/extmod/btree_dup.py @@ -0,0 +1,160 @@ +import btree +import ustruct +from ucollections import namedtuple + +#this class is intended to be in micropython-lib/dbm module +class dbm: + + def __init__(self,file,psize=512,keyformat='>i', valformat='>i',dup=False): + self.db=None + self.f=None + self.file = file + self.psize = psize + self.kb=bytearray(ustruct.calcsize(keyformat)) + self.vb=bytearray(ustruct.calcsize(valformat)) + self.keyformat=keyformat + self.valformat=valformat + self.vnt=None + self.knt=None + self.dup=dup + + def __enter__(self): + try: + self.f = open(self.file,'r+b') + self.db = btree.open(self.f,pagesize=self.psize,flags=btree.DUPLICATE_KEY if self.dup is True else 0) + return self + except OSError: + try: + self.f = open(self.file,'w+b') + self.db = btree.open(self.f,pagesize=self.psize,flags=btree.DUPLICATE_KEY if self.dup is True else 0) + return self + except OSError as e: + print(repr(e)) + if self.f is not None: + self.f.close() + + def __exit__(self,exc_type,value,traceback): + if self.db is not None: + self.db.close() + if self.f is not None: + self.f.close() + + @staticmethod + def pack_data(data,_format,buf): + if isinstance(data,tuple): + ustruct.pack_into(_format,buf,0,*data) + else: + ustruct.pack_into(_format,buf,0,data) + + def get_size(self): + self.f.seek(0,2) + size = self.f.tell() + return size + + def __setitem__(self,key,val): + dbm.pack_data(key,self.keyformat,self.kb) + dbm.pack_data(val,self.valformat,self.vb) + self.db[bytes(self.kb)]=bytes(self.vb) + + def ntuple(self,nt,data): + if nt is not None: + return nt(*data) + else: + return data + + def _iter(self,items): + for i in items: + yield self.ntuple(self.knt,ustruct.unpack_from(self.keyformat,i[0])),self.ntuple(self.vnt,ustruct.unpack_from(self.valformat,i[1])) + + def __getitem__(self,key): + if isinstance(key,slice): + dbm.pack_data(key.start,self.keyformat,self.kb) + start=bytes(self.kb) + dbm.pack_data(key.stop,self.keyformat,self.kb) + stop=bytes(self.kb) + items = self.db.items(start,stop) + return iter(self._iter(items)) + else: + dbm.pack_data(key,self.keyformat,self.kb) + val = self.db[bytes(self.kb)] + return self.ntuple(self.vnt,ustruct.unpack_from(self.valformat,val)) + + def __iter__(self): + return iter(self._iter(self.db.items())) + + def __delitem__(self,key): + dbm.pack_data(key,self.keyformat,self.kb) + del self.db[bytes(self.kb)] + + def set_key_ntuple(self,*field_names): + self.knt = namedtuple('Key',field_names) + return self.knt + + def set_val_ntuple(self,*field_names): + self.vnt = namedtuple('Val',field_names) + return self.vnt + + +def set_db(j): + with dbm(f_name,keyformat=k_fmt,valformat=v_fmt,dup=dupplic) as d: + for i in range(j): + d[ i//dup_div+1, (i//dup_div)%100, 100-((i//dup_div)%100), ((i//dup_div)%100)*1010] = 20170000+100*((i//100)%12+1)+((i//100)%31)+1, i, (i//100)%2, 1, i%1000, 10, (i//10)%2 + +def read_db(): + with dbm(f_name,keyformat=k_fmt,valformat=v_fmt,dup=dupplic) as d: + d.set_key_ntuple(*k_t) + d.set_val_ntuple(*v_t) + for i in d: + print ("read:", i) + +def read_at(k11=0, k12=0, k13=0, k14=0, k21=99999, k22=255, k23=99999, k24=9999999): + with dbm(f_name,keyformat=k_fmt,valformat=v_fmt,dup=dupplic) as d: + d.set_key_ntuple(*k_t) + d.set_val_ntuple(*v_t) + ii=1 + for i in d[(k11,):(k21,)]: + if i[0].k2>=k12 and i[0].k2<=k22 and i[0].k3>=k13 and i[0].k3<=k23 and i[0].k4>=k14 and i[0].k4<=k24: + print ("sel: ",ii, i) + ii+=1 + +def delete_at (): + with dbm(f_name,keyformat=k_fmt,valformat=v_fmt,dup=dupplic) as d: + d.set_key_ntuple(*k_t) + d.set_val_ntuple(*v_t) + ii=1 + t=(7,6,94,6060) + print ("del: ",ii,t, d[t]) + del d[t] + ii+=1 + +#----------------------------------------------------------------- +# main: + +f_name = "file.db" +k_fmt = "!HBHI" +k_t = ('k1', 'k2', 'k3', 'k4') +v_fmt = "!IIBBHHB" +v_t = ('v1', 'v2', 'v3', 'v4', 'v5', 'v6', 'v7') +dupplic=True # True => set dbm to manage dupplicate keys +dup_div= 10 + +with open(f_name, "wb") as f: + pass + +set_db(100) +read_db() +read_at(k11=7, k21=8) + +set_db(100) +read_db() +read_at(k11=7, k21=8) + +set_db(100) +read_db() +read_at(k11=7, k21=8) + +set_db(100) +read_db() +read_at(k11=7, k21=8) + +delete_at()