diff --git a/pysondb/db.py b/pysondb/db.py index e615e34..0506211 100644 --- a/pysondb/db.py +++ b/pysondb/db.py @@ -99,7 +99,7 @@ def add(self, data: object) -> str: keys = db_data['keys'] if not isinstance(keys, list): raise SchemaTypeError( - f"keys must of type 'list' and not {type(keys)}") + f"keys must be of type 'list' and not {type(keys)}") if len(keys) == 0: db_data['keys'] = sorted(list(data.keys())) else: @@ -164,6 +164,23 @@ def add_many(self, data: object, json_response: bool = False) -> Union[SingleDat return new_data if json_response else None + def get_all_select_keys(self, keys: List[str]) -> ReturnWithIdType: + with self.lock: + file_data = self._load_file() + data = file_data['data'] + existing_keys = file_data['keys'] + if not isinstance(existing_keys, list): + raise SchemaTypeError( + f"keys must be of type 'list' and not {type(existing_keys)}") + if not set(keys).issubset(existing_keys): + raise UnknownKeyError( + f'Unrecognized key(s) {set(keys) ^ set(existing_keys)}' + '(Unable to find the key(s) in the DB)' + ) + if isinstance(data, dict): + return {k: {i: v[i] for i in keys} for k, v in data.items()} + return {} + def get_all(self) -> ReturnWithIdType: with self.lock: data = self._load_file()['data'] diff --git a/setup.cfg b/setup.cfg index 3cba140..7eeff78 100644 --- a/setup.cfg +++ b/setup.cfg @@ -1,6 +1,6 @@ [metadata] name = pysondb-v2 -version = 2.1.0 +version = 2.2.0 author = Adwaith Rajesh author_email = adwaithrajesh3180@gmail.com description = A Simple, Lightweight, Efficent JSON based database for Python. diff --git a/tests/test_db_get_all_select_keys.py b/tests/test_db_get_all_select_keys.py new file mode 100644 index 0000000..3378897 --- /dev/null +++ b/tests/test_db_get_all_select_keys.py @@ -0,0 +1,63 @@ +import json +import pytest +from pysondb.db import PysonDB +from pysondb.errors import UnknownKeyError + + +TEST_DATA = { + 'version': 2, + 'keys': ['age', 'name', 'toy'], + 'data': { + '2352346': { + 'age': 4, + 'name': 'mathew_first', + 'toy': 'car', + }, + '1234567': { + 'age': 9, + 'name': 'new_user', + 'toy': 'ball', + } + } +} + + +def test_get_all_select_keys(tmpdir): + f = tmpdir.join('test.json') + f.write(json.dumps(TEST_DATA)) + + db = PysonDB(f.strpath) + assert db.get_all_select_keys(['age']) == { + '2352346': { + 'age': 4, + }, + '1234567': { + 'age': 9, + } + } + assert db.get_all_select_keys(['name', 'toy']) == { + '2352346': { + 'name': 'mathew_first', + 'toy': 'car', + }, + '1234567': { + 'name': 'new_user', + 'toy': 'ball', + } + } + + +def test_get_all_select_keys_empty_file(tmpdir): + f = tmpdir.join('test.json') + f.write(json.dumps({'version': 2, 'keys': [], 'data': {}})) + + db = PysonDB(f.strpath) + assert db.get_all_select_keys([]) == {} + +def test_get_all_select_keys_wrong_key(tmpdir): + f = tmpdir.join('test.json') + f.write(json.dumps(TEST_DATA)) + + db = PysonDB(f.strpath) + with pytest.raises(UnknownKeyError): + db.get_all_select_keys(['wrong_key'])