Return Plain Dicts from BasePersistence.get_*_data (#2873)

This commit is contained in:
Harshil 2022-02-01 20:50:17 +04:00 committed by Hinrich Mahler
parent 778eadbe7c
commit 835434c12f
5 changed files with 119 additions and 150 deletions

View file

@ -19,7 +19,7 @@
"""This module contains the BasePersistence class.""" """This module contains the BasePersistence class."""
from abc import ABC, abstractmethod from abc import ABC, abstractmethod
from copy import copy from copy import copy
from typing import Dict, Optional, Tuple, cast, ClassVar, Generic, DefaultDict, NamedTuple from typing import Dict, Optional, Tuple, cast, ClassVar, Generic, NamedTuple
from telegram import Bot from telegram import Bot
from telegram.ext import ExtBot from telegram.ext import ExtBot
@ -132,10 +132,10 @@ class BasePersistence(Generic[UD, CD, BD], ABC):
update_bot_data = instance.update_bot_data update_bot_data = instance.update_bot_data
update_callback_data = instance.update_callback_data update_callback_data = instance.update_callback_data
def get_user_data_insert_bot() -> DefaultDict[int, UD]: def get_user_data_insert_bot() -> Dict[int, UD]:
return instance.insert_bot(get_user_data()) return instance.insert_bot(get_user_data())
def get_chat_data_insert_bot() -> DefaultDict[int, CD]: def get_chat_data_insert_bot() -> Dict[int, CD]:
return instance.insert_bot(get_chat_data()) return instance.insert_bot(get_chat_data())
def get_bot_data_insert_bot() -> BD: def get_bot_data_insert_bot() -> BD:
@ -171,10 +171,7 @@ class BasePersistence(Generic[UD, CD, BD], ABC):
setattr(instance, 'update_callback_data', update_callback_data_replace_bot) setattr(instance, 'update_callback_data', update_callback_data_replace_bot)
return instance return instance
def __init__( def __init__(self, store_data: PersistenceInput = None):
self,
store_data: PersistenceInput = None,
):
self.store_data = store_data or PersistenceInput() self.store_data = store_data or PersistenceInput()
self.bot: Bot = None # type: ignore[assignment] self.bot: Bot = None # type: ignore[assignment]
@ -398,34 +395,40 @@ class BasePersistence(Generic[UD, CD, BD], ABC):
return obj return obj
@abstractmethod @abstractmethod
def get_user_data(self) -> DefaultDict[int, UD]: def get_user_data(self) -> Dict[int, UD]:
"""Will be called by :class:`telegram.ext.Dispatcher` upon creation with a """Will be called by :class:`telegram.ext.Dispatcher` upon creation with a
persistence object. It should return the ``user_data`` if stored, or an empty persistence object. It should return the ``user_data`` if stored, or an empty
:obj:`defaultdict`. In the latter case, the :obj:`defaultdict` should produce values :obj:`dict`. In the latter case, the dictionary should produce values
corresponding to one of the following: corresponding to one of the following:
* :obj:`dict` * :obj:`dict`
* The type from :attr:`telegram.ext.ContextTypes.user_data` * The type from :attr:`telegram.ext.ContextTypes.user_data`
if :class:`telegram.ext.ContextTypes` are used. if :class:`telegram.ext.ContextTypes` is used.
.. versionchanged:: 14.0
This method may now return a :obj:`dict` instead of a :obj:`collections.defaultdict`
Returns: Returns:
DefaultDict[:obj:`int`, :obj:`dict` | :attr:`telegram.ext.ContextTypes.user_data`]: Dict[:obj:`int`, :obj:`dict` | :attr:`telegram.ext.ContextTypes.user_data`]:
The restored user data. The restored user data.
""" """
@abstractmethod @abstractmethod
def get_chat_data(self) -> DefaultDict[int, CD]: def get_chat_data(self) -> Dict[int, CD]:
"""Will be called by :class:`telegram.ext.Dispatcher` upon creation with a """Will be called by :class:`telegram.ext.Dispatcher` upon creation with a
persistence object. It should return the ``chat_data`` if stored, or an empty persistence object. It should return the ``chat_data`` if stored, or an empty
:obj:`defaultdict`. In the latter case, the :obj:`defaultdict` should produce values :obj:`dict`. In the latter case, the dictionary should produce values
corresponding to one of the following: corresponding to one of the following:
* :obj:`dict` * :obj:`dict`
* The type from :attr:`telegram.ext.ContextTypes.chat_data` * The type from :attr:`telegram.ext.ContextTypes.chat_data`
if :class:`telegram.ext.ContextTypes` are used. if :class:`telegram.ext.ContextTypes` is used.
.. versionchanged:: 14.0
This method may now return a :obj:`dict` instead of a :obj:`collections.defaultdict`
Returns: Returns:
DefaultDict[:obj:`int`, :obj:`dict` | :attr:`telegram.ext.ContextTypes.chat_data`]: Dict[:obj:`int`, :obj:`dict` | :attr:`telegram.ext.ContextTypes.chat_data`]:
The restored chat data. The restored chat data.
""" """
@ -433,7 +436,7 @@ class BasePersistence(Generic[UD, CD, BD], ABC):
def get_bot_data(self) -> BD: def get_bot_data(self) -> BD:
"""Will be called by :class:`telegram.ext.Dispatcher` upon creation with a """Will be called by :class:`telegram.ext.Dispatcher` upon creation with a
persistence object. It should return the ``bot_data`` if stored, or an empty persistence object. It should return the ``bot_data`` if stored, or an empty
:obj:`defaultdict`. In the latter case, the :obj:`defaultdict` should produce values :obj:`dict`. In the latter case, the :obj:`dict` should produce values
corresponding to one of the following: corresponding to one of the following:
* :obj:`dict` * :obj:`dict`
@ -441,7 +444,7 @@ class BasePersistence(Generic[UD, CD, BD], ABC):
if :class:`telegram.ext.ContextTypes` are used. if :class:`telegram.ext.ContextTypes` are used.
Returns: Returns:
DefaultDict[:obj:`int`, :obj:`dict` | :attr:`telegram.ext.ContextTypes.bot_data`]: Dict[:obj:`int`, :obj:`dict` | :attr:`telegram.ext.ContextTypes.bot_data`]:
The restored bot data. The restored bot data.
""" """

View file

@ -18,8 +18,7 @@
# along with this program. If not, see [http://www.gnu.org/licenses/]. # along with this program. If not, see [http://www.gnu.org/licenses/].
"""This module contains the DictPersistence class.""" """This module contains the DictPersistence class."""
from typing import DefaultDict, Dict, Optional, Tuple, cast from typing import Dict, Optional, Tuple, cast
from collections import defaultdict
from telegram.ext import BasePersistence, PersistenceInput from telegram.ext import BasePersistence, PersistenceInput
from telegram._utils.types import JSONDict from telegram._utils.types import JSONDict
@ -166,7 +165,7 @@ class DictPersistence(BasePersistence):
) from exc ) from exc
@property @property
def user_data(self) -> Optional[DefaultDict[int, Dict]]: def user_data(self) -> Optional[Dict[int, Dict]]:
""":obj:`dict`: The user_data as a dict.""" """:obj:`dict`: The user_data as a dict."""
return self._user_data return self._user_data
@ -178,7 +177,7 @@ class DictPersistence(BasePersistence):
return json.dumps(self.user_data) return json.dumps(self.user_data)
@property @property
def chat_data(self) -> Optional[DefaultDict[int, Dict]]: def chat_data(self) -> Optional[Dict[int, Dict]]:
""":obj:`dict`: The chat_data as a dict.""" """:obj:`dict`: The chat_data as a dict."""
return self._chat_data return self._chat_data
@ -204,7 +203,7 @@ class DictPersistence(BasePersistence):
@property @property
def callback_data(self) -> Optional[CDCData]: def callback_data(self) -> Optional[CDCData]:
"""Tuple[List[Tuple[:obj:`str`, :obj:`float`, Dict[:obj:`str`, :obj:`Any`]]], \ """Tuple[List[Tuple[:obj:`str`, :obj:`float`, Dict[:obj:`str`, :obj:`Any`]]], \
Dict[:obj:`str`, :obj:`str`]]: The meta data on the stored callback data. Dict[:obj:`str`, :obj:`str`]]: The metadata on the stored callback data.
.. versionadded:: 13.6 .. versionadded:: 13.6
""" """
@ -212,7 +211,7 @@ class DictPersistence(BasePersistence):
@property @property
def callback_data_json(self) -> str: def callback_data_json(self) -> str:
""":obj:`str`: The meta data on the stored callback data as a JSON-string. """:obj:`str`: The metadata on the stored callback data as a JSON-string.
.. versionadded:: 13.6 .. versionadded:: 13.6
""" """
@ -232,26 +231,24 @@ class DictPersistence(BasePersistence):
return self._conversations_json return self._conversations_json
return self._encode_conversations_to_json(self.conversations) # type: ignore[arg-type] return self._encode_conversations_to_json(self.conversations) # type: ignore[arg-type]
def get_user_data(self) -> DefaultDict[int, Dict[object, object]]: def get_user_data(self) -> Dict[int, Dict[object, object]]:
"""Returns the user_data created from the ``user_data_json`` or an empty """Returns the user_data created from the ``user_data_json`` or an empty :obj:`dict`.
:obj:`defaultdict`.
Returns: Returns:
:obj:`defaultdict`: The restored user data. :obj:`dict`: The restored user data.
""" """
if self.user_data is None: if self.user_data is None:
self._user_data = defaultdict(dict) self._user_data = {}
return self.user_data # type: ignore[return-value] return self.user_data # type: ignore[return-value]
def get_chat_data(self) -> DefaultDict[int, Dict[object, object]]: def get_chat_data(self) -> Dict[int, Dict[object, object]]:
"""Returns the chat_data created from the ``chat_data_json`` or an empty """Returns the chat_data created from the ``chat_data_json`` or an empty :obj:`dict`.
:obj:`defaultdict`.
Returns: Returns:
:obj:`defaultdict`: The restored chat data. :obj:`dict`: The restored chat data.
""" """
if self.chat_data is None: if self.chat_data is None:
self._chat_data = defaultdict(dict) self._chat_data = {}
return self.chat_data # type: ignore[return-value] return self.chat_data # type: ignore[return-value]
def get_bot_data(self) -> Dict[object, object]: def get_bot_data(self) -> Dict[object, object]:
@ -271,7 +268,7 @@ class DictPersistence(BasePersistence):
Returns: Returns:
Tuple[List[Tuple[:obj:`str`, :obj:`float`, Dict[:obj:`str`, :obj:`Any`]]], \ Tuple[List[Tuple[:obj:`str`, :obj:`float`, Dict[:obj:`str`, :obj:`Any`]]], \
Dict[:obj:`str`, :obj:`str`]]: The restored meta data or :obj:`None`, \ Dict[:obj:`str`, :obj:`str`]]: The restored metadata or :obj:`None`, \
if no data was stored. if no data was stored.
""" """
if self.callback_data is None: if self.callback_data is None:
@ -315,7 +312,7 @@ class DictPersistence(BasePersistence):
data (:obj:`dict`): The :attr:`telegram.ext.Dispatcher.user_data` ``[user_id]``. data (:obj:`dict`): The :attr:`telegram.ext.Dispatcher.user_data` ``[user_id]``.
""" """
if self._user_data is None: if self._user_data is None:
self._user_data = defaultdict(dict) self._user_data = {}
if self._user_data.get(user_id) == data: if self._user_data.get(user_id) == data:
return return
self._user_data[user_id] = data self._user_data[user_id] = data
@ -329,7 +326,7 @@ class DictPersistence(BasePersistence):
data (:obj:`dict`): The :attr:`telegram.ext.Dispatcher.chat_data` ``[chat_id]``. data (:obj:`dict`): The :attr:`telegram.ext.Dispatcher.chat_data` ``[chat_id]``.
""" """
if self._chat_data is None: if self._chat_data is None:
self._chat_data = defaultdict(dict) self._chat_data = {}
if self._chat_data.get(chat_id) == data: if self._chat_data.get(chat_id) == data:
return return
self._chat_data[chat_id] = data self._chat_data[chat_id] = data
@ -453,7 +450,7 @@ class DictPersistence(BasePersistence):
return conversations return conversations
@staticmethod @staticmethod
def _decode_user_chat_data_from_json(data: str) -> DefaultDict[int, Dict[object, object]]: def _decode_user_chat_data_from_json(data: str) -> Dict[int, Dict[object, object]]:
"""Helper method to decode chat or user data (that uses ints as keys) from a """Helper method to decode chat or user data (that uses ints as keys) from a
JSON-string. JSON-string.
@ -463,7 +460,7 @@ class DictPersistence(BasePersistence):
Returns: Returns:
:obj:`dict`: The user/chat_data defaultdict after decoding :obj:`dict`: The user/chat_data defaultdict after decoding
""" """
tmp: DefaultDict[int, Dict[object, object]] = defaultdict(dict) tmp: Dict[int, Dict[object, object]] = {}
decoded_data = json.loads(data) decoded_data = json.loads(data)
for user, user_data in decoded_data.items(): for user, user_data in decoded_data.items():
user = int(user) user = int(user)

View file

@ -739,7 +739,7 @@ class Dispatcher(Generic[BT, CCT, UD, CD, BD, JQ, PT]):
def __update_persistence(self, update: object = None) -> None: def __update_persistence(self, update: object = None) -> None:
if self.persistence: if self.persistence:
# We use list() here in order to decouple chat_ids from self._chat_data, as dict view # We use list() here in order to decouple chat_ids from self.chat_data, as dict view
# objects will change, when the dict does and we want to loop over chat_ids # objects will change, when the dict does and we want to loop over chat_ids
chat_ids = list(self.chat_data.keys()) chat_ids = list(self.chat_data.keys())
user_ids = list(self.user_data.keys()) user_ids = list(self.user_data.keys())

View file

@ -18,7 +18,6 @@
# along with this program. If not, see [http://www.gnu.org/licenses/]. # along with this program. If not, see [http://www.gnu.org/licenses/].
"""This module contains the PicklePersistence class.""" """This module contains the PicklePersistence class."""
import pickle import pickle
from collections import defaultdict
from pathlib import Path from pathlib import Path
from typing import ( from typing import (
Any, Any,
@ -27,7 +26,6 @@ from typing import (
Tuple, Tuple,
overload, overload,
cast, cast,
DefaultDict,
) )
from telegram._utils.types import FilePathInput from telegram._utils.types import FilePathInput
@ -138,8 +136,8 @@ class PicklePersistence(BasePersistence[UD, CD, BD]):
self.filepath = Path(filepath) self.filepath = Path(filepath)
self.single_file = single_file self.single_file = single_file
self.on_flush = on_flush self.on_flush = on_flush
self.user_data: Optional[DefaultDict[int, UD]] = None self.user_data: Optional[Dict[int, UD]] = None
self.chat_data: Optional[DefaultDict[int, CD]] = None self.chat_data: Optional[Dict[int, CD]] = None
self.bot_data: Optional[BD] = None self.bot_data: Optional[BD] = None
self.callback_data: Optional[CDCData] = None self.callback_data: Optional[CDCData] = None
self.conversations: Optional[Dict[str, Dict[Tuple, object]]] = None self.conversations: Optional[Dict[str, Dict[Tuple, object]]] = None
@ -149,16 +147,16 @@ class PicklePersistence(BasePersistence[UD, CD, BD]):
try: try:
with self.filepath.open("rb") as file: with self.filepath.open("rb") as file:
data = pickle.load(file) data = pickle.load(file)
self.user_data = defaultdict(self.context_types.user_data, data['user_data']) self.user_data = data['user_data']
self.chat_data = defaultdict(self.context_types.chat_data, data['chat_data']) self.chat_data = data['chat_data']
# For backwards compatibility with files not containing bot data # For backwards compatibility with files not containing bot data
self.bot_data = data.get('bot_data', self.context_types.bot_data()) self.bot_data = data.get('bot_data', self.context_types.bot_data())
self.callback_data = data.get('callback_data', {}) self.callback_data = data.get('callback_data', {})
self.conversations = data['conversations'] self.conversations = data['conversations']
except OSError: except OSError:
self.conversations = {} self.conversations = {}
self.user_data = defaultdict(self.context_types.user_data) self.user_data = {}
self.chat_data = defaultdict(self.context_types.chat_data) self.chat_data = {}
self.bot_data = self.context_types.bot_data() self.bot_data = self.context_types.bot_data()
self.callback_data = None self.callback_data = None
except pickle.UnpicklingError as exc: except pickle.UnpicklingError as exc:
@ -195,41 +193,35 @@ class PicklePersistence(BasePersistence[UD, CD, BD]):
with filepath.open("wb") as file: with filepath.open("wb") as file:
pickle.dump(data, file) pickle.dump(data, file)
def get_user_data(self) -> DefaultDict[int, UD]: def get_user_data(self) -> Dict[int, UD]:
"""Returns the user_data from the pickle file if it exists or an empty :obj:`defaultdict`. """Returns the user_data from the pickle file if it exists or an empty :obj:`dict`.
Returns: Returns:
DefaultDict[:obj:`int`, :obj:`dict` | :attr:`telegram.ext.ContextTypes.user_data`]: Dict[:obj:`int`, :obj:`dict`]: The restored user data.
The restored user data.
""" """
if self.user_data: if self.user_data:
pass pass
elif not self.single_file: elif not self.single_file:
data = self._load_file(Path(f"{self.filepath}_user_data")) data = self._load_file(Path(f"{self.filepath}_user_data"))
if not data: if not data:
data = defaultdict(self.context_types.user_data) data = {}
else:
data = defaultdict(self.context_types.user_data, data)
self.user_data = data self.user_data = data
else: else:
self._load_singlefile() self._load_singlefile()
return self.user_data # type: ignore[return-value] return self.user_data # type: ignore[return-value]
def get_chat_data(self) -> DefaultDict[int, CD]: def get_chat_data(self) -> Dict[int, CD]:
"""Returns the chat_data from the pickle file if it exists or an empty :obj:`defaultdict`. """Returns the chat_data from the pickle file if it exists or an empty :obj:`dict`.
Returns: Returns:
DefaultDict[:obj:`int`, :obj:`dict` | :attr:`telegram.ext.ContextTypes.chat_data`]: Dict[:obj:`int`, :obj:`dict`]: The restored chat data.
The restored chat data.
""" """
if self.chat_data: if self.chat_data:
pass pass
elif not self.single_file: elif not self.single_file:
data = self._load_file(Path(f"{self.filepath}_chat_data")) data = self._load_file(Path(f"{self.filepath}_chat_data"))
if not data: if not data:
data = defaultdict(self.context_types.chat_data) data = {}
else:
data = defaultdict(self.context_types.chat_data, data)
self.chat_data = data self.chat_data = data
else: else:
self._load_singlefile() self._load_singlefile()
@ -260,8 +252,8 @@ class PicklePersistence(BasePersistence[UD, CD, BD]):
Returns: Returns:
Optional[Tuple[List[Tuple[:obj:`str`, :obj:`float`, \ Optional[Tuple[List[Tuple[:obj:`str`, :obj:`float`, \
Dict[:obj:`str`, :obj:`Any`]]], Dict[:obj:`str`, :obj:`str`]]: Dict[:obj:`str`, :obj:`Any`]]], Dict[:obj:`str`, :obj:`str`]]]:
The restored meta data or :obj:`None`, if no data was stored. The restored metadata or :obj:`None`, if no data was stored.
""" """
if self.callback_data: if self.callback_data:
pass pass
@ -323,11 +315,10 @@ class PicklePersistence(BasePersistence[UD, CD, BD]):
Args: Args:
user_id (:obj:`int`): The user the data might have been changed for. user_id (:obj:`int`): The user the data might have been changed for.
data (:obj:`dict` | :attr:`telegram.ext.ContextTypes.user_data`): The data (:obj:`dict`): The :attr:`telegram.ext.Dispatcher.user_data` ``[user_id]``.
:attr:`telegram.ext.Dispatcher.user_data` ``[user_id]``.
""" """
if self.user_data is None: if self.user_data is None:
self.user_data = defaultdict(self.context_types.user_data) self.user_data = {}
if self.user_data.get(user_id) == data: if self.user_data.get(user_id) == data:
return return
self.user_data[user_id] = data self.user_data[user_id] = data
@ -342,11 +333,10 @@ class PicklePersistence(BasePersistence[UD, CD, BD]):
Args: Args:
chat_id (:obj:`int`): The chat the data might have been changed for. chat_id (:obj:`int`): The chat the data might have been changed for.
data (:obj:`dict` | :attr:`telegram.ext.ContextTypes.chat_data`): The data (:obj:`dict`): The :attr:`telegram.ext.Dispatcher.chat_data` ``[chat_id]``.
:attr:`telegram.ext.Dispatcher.chat_data` ``[chat_id]``.
""" """
if self.chat_data is None: if self.chat_data is None:
self.chat_data = defaultdict(self.context_types.chat_data) self.chat_data = {}
if self.chat_data.get(chat_id) == data: if self.chat_data.get(chat_id) == data:
return return
self.chat_data[chat_id] = data self.chat_data[chat_id] = data

View file

@ -16,26 +16,24 @@
# #
# You should have received a copy of the GNU Lesser Public License # You should have received a copy of the GNU Lesser Public License
# along with this program. If not, see [http://www.gnu.org/licenses/]. # along with this program. If not, see [http://www.gnu.org/licenses/].
import logging
import os
import pickle
import gzip import gzip
import signal import signal
import uuid import uuid
from collections.abc import Container
from collections import defaultdict
from pathlib import Path from pathlib import Path
from time import sleep
from threading import Lock from threading import Lock
from telegram.ext import PersistenceInput, UpdaterBuilder, CallbackDataCache import pytest
try: try:
import ujson as json import ujson as json
except ImportError: except ImportError:
import json import json
import logging
import os
import pickle
from collections import defaultdict
from collections.abc import Container
from time import sleep
import pytest
from telegram import Update, Message, User, Chat, MessageEntity, Bot from telegram import Update, Message, User, Chat, MessageEntity, Bot
from telegram.ext import ( from telegram.ext import (
@ -49,6 +47,9 @@ from telegram.ext import (
TypeHandler, TypeHandler,
JobQueue, JobQueue,
ContextTypes, ContextTypes,
PersistenceInput,
UpdaterBuilder,
CallbackDataCache,
) )
from telegram.ext._callbackdatacache import _KeyboardData from telegram.ext._callbackdatacache import _KeyboardData
@ -134,8 +135,8 @@ def bot_persistence():
def __init__(self): def __init__(self):
super().__init__() super().__init__()
self.bot_data = None self.bot_data = None
self.chat_data = defaultdict(dict) self.chat_data = {}
self.user_data = defaultdict(dict) self.user_data = {}
self.callback_data = None self.callback_data = None
def get_bot_data(self): def get_bot_data(self):
@ -196,16 +197,12 @@ def bot_data():
@pytest.fixture(scope="function") @pytest.fixture(scope="function")
def chat_data(): def chat_data():
return defaultdict( return {-12345: {'test1': 'test2', 'test3': {'test4': 'test5'}}, -67890: {3: 'test4'}}
dict, {-12345: {'test1': 'test2', 'test3': {'test4': 'test5'}}, -67890: {3: 'test4'}}
)
@pytest.fixture(scope="function") @pytest.fixture(scope="function")
def user_data(): def user_data():
return defaultdict( return {12345: {'test1': 'test2', 'test3': {'test4': 'test5'}}, 67890: {3: 'test4'}}
dict, {12345: {'test1': 'test2', 'test3': {'test4': 'test5'}}, 67890: {3: 'test4'}}
)
@pytest.fixture(scope="function") @pytest.fixture(scope="function")
@ -1040,10 +1037,6 @@ def update(bot):
return Update(0, message=message) return Update(0, message=message)
class CustomMapping(defaultdict):
pass
class TestPicklePersistence: class TestPicklePersistence:
def test_slot_behaviour(self, mro_slots, pickle_persistence): def test_slot_behaviour(self, mro_slots, pickle_persistence):
inst = pickle_persistence inst = pickle_persistence
@ -1059,16 +1052,16 @@ class TestPicklePersistence:
assert retrieved == bot_data assert retrieved == bot_data
def test_no_files_present_multi_file(self, pickle_persistence): def test_no_files_present_multi_file(self, pickle_persistence):
assert pickle_persistence.get_user_data() == defaultdict(dict) assert pickle_persistence.get_user_data() == {}
assert pickle_persistence.get_chat_data() == defaultdict(dict) assert pickle_persistence.get_chat_data() == {}
assert pickle_persistence.get_bot_data() == {} assert pickle_persistence.get_bot_data() == {}
assert pickle_persistence.get_callback_data() is None assert pickle_persistence.get_callback_data() is None
assert pickle_persistence.get_conversations('noname') == {} assert pickle_persistence.get_conversations('noname') == {}
def test_no_files_present_single_file(self, pickle_persistence): def test_no_files_present_single_file(self, pickle_persistence):
pickle_persistence.single_file = True pickle_persistence.single_file = True
assert pickle_persistence.get_user_data() == defaultdict(dict) assert pickle_persistence.get_user_data() == {}
assert pickle_persistence.get_chat_data() == defaultdict(dict) assert pickle_persistence.get_chat_data() == {}
assert pickle_persistence.get_bot_data() == {} assert pickle_persistence.get_bot_data() == {}
assert pickle_persistence.get_callback_data() is None assert pickle_persistence.get_callback_data() is None
assert pickle_persistence.get_conversations('noname') == {} assert pickle_persistence.get_conversations('noname') == {}
@ -1125,16 +1118,14 @@ class TestPicklePersistence:
def test_with_good_multi_file(self, pickle_persistence, good_pickle_files): def test_with_good_multi_file(self, pickle_persistence, good_pickle_files):
user_data = pickle_persistence.get_user_data() user_data = pickle_persistence.get_user_data()
assert isinstance(user_data, defaultdict) assert isinstance(user_data, dict)
assert user_data[12345]['test1'] == 'test2' assert user_data[12345]['test1'] == 'test2'
assert user_data[67890][3] == 'test4' assert user_data[67890][3] == 'test4'
assert user_data[54321] == {}
chat_data = pickle_persistence.get_chat_data() chat_data = pickle_persistence.get_chat_data()
assert isinstance(chat_data, defaultdict) assert isinstance(chat_data, dict)
assert chat_data[-12345]['test1'] == 'test2' assert chat_data[-12345]['test1'] == 'test2'
assert chat_data[-67890][3] == 'test4' assert chat_data[-67890][3] == 'test4'
assert chat_data[-54321] == {}
bot_data = pickle_persistence.get_bot_data() bot_data = pickle_persistence.get_bot_data()
assert isinstance(bot_data, dict) assert isinstance(bot_data, dict)
@ -1163,16 +1154,14 @@ class TestPicklePersistence:
def test_with_good_single_file(self, pickle_persistence, good_pickle_files): def test_with_good_single_file(self, pickle_persistence, good_pickle_files):
pickle_persistence.single_file = True pickle_persistence.single_file = True
user_data = pickle_persistence.get_user_data() user_data = pickle_persistence.get_user_data()
assert isinstance(user_data, defaultdict) assert isinstance(user_data, dict)
assert user_data[12345]['test1'] == 'test2' assert user_data[12345]['test1'] == 'test2'
assert user_data[67890][3] == 'test4' assert user_data[67890][3] == 'test4'
assert user_data[54321] == {}
chat_data = pickle_persistence.get_chat_data() chat_data = pickle_persistence.get_chat_data()
assert isinstance(chat_data, defaultdict) assert isinstance(chat_data, dict)
assert chat_data[-12345]['test1'] == 'test2' assert chat_data[-12345]['test1'] == 'test2'
assert chat_data[-67890][3] == 'test4' assert chat_data[-67890][3] == 'test4'
assert chat_data[-54321] == {}
bot_data = pickle_persistence.get_bot_data() bot_data = pickle_persistence.get_bot_data()
assert isinstance(bot_data, dict) assert isinstance(bot_data, dict)
@ -1200,16 +1189,14 @@ class TestPicklePersistence:
def test_with_multi_file_wo_bot_data(self, pickle_persistence, pickle_files_wo_bot_data): def test_with_multi_file_wo_bot_data(self, pickle_persistence, pickle_files_wo_bot_data):
user_data = pickle_persistence.get_user_data() user_data = pickle_persistence.get_user_data()
assert isinstance(user_data, defaultdict) assert isinstance(user_data, dict)
assert user_data[12345]['test1'] == 'test2' assert user_data[12345]['test1'] == 'test2'
assert user_data[67890][3] == 'test4' assert user_data[67890][3] == 'test4'
assert user_data[54321] == {}
chat_data = pickle_persistence.get_chat_data() chat_data = pickle_persistence.get_chat_data()
assert isinstance(chat_data, defaultdict) assert isinstance(chat_data, dict)
assert chat_data[-12345]['test1'] == 'test2' assert chat_data[-12345]['test1'] == 'test2'
assert chat_data[-67890][3] == 'test4' assert chat_data[-67890][3] == 'test4'
assert chat_data[-54321] == {}
bot_data = pickle_persistence.get_bot_data() bot_data = pickle_persistence.get_bot_data()
assert isinstance(bot_data, dict) assert isinstance(bot_data, dict)
@ -1237,16 +1224,14 @@ class TestPicklePersistence:
self, pickle_persistence, pickle_files_wo_callback_data self, pickle_persistence, pickle_files_wo_callback_data
): ):
user_data = pickle_persistence.get_user_data() user_data = pickle_persistence.get_user_data()
assert isinstance(user_data, defaultdict) assert isinstance(user_data, dict)
assert user_data[12345]['test1'] == 'test2' assert user_data[12345]['test1'] == 'test2'
assert user_data[67890][3] == 'test4' assert user_data[67890][3] == 'test4'
assert user_data[54321] == {}
chat_data = pickle_persistence.get_chat_data() chat_data = pickle_persistence.get_chat_data()
assert isinstance(chat_data, defaultdict) assert isinstance(chat_data, dict)
assert chat_data[-12345]['test1'] == 'test2' assert chat_data[-12345]['test1'] == 'test2'
assert chat_data[-67890][3] == 'test4' assert chat_data[-67890][3] == 'test4'
assert chat_data[-54321] == {}
bot_data = pickle_persistence.get_bot_data() bot_data = pickle_persistence.get_bot_data()
assert isinstance(bot_data, dict) assert isinstance(bot_data, dict)
@ -1273,16 +1258,14 @@ class TestPicklePersistence:
def test_with_single_file_wo_bot_data(self, pickle_persistence, pickle_files_wo_bot_data): def test_with_single_file_wo_bot_data(self, pickle_persistence, pickle_files_wo_bot_data):
pickle_persistence.single_file = True pickle_persistence.single_file = True
user_data = pickle_persistence.get_user_data() user_data = pickle_persistence.get_user_data()
assert isinstance(user_data, defaultdict) assert isinstance(user_data, dict)
assert user_data[12345]['test1'] == 'test2' assert user_data[12345]['test1'] == 'test2'
assert user_data[67890][3] == 'test4' assert user_data[67890][3] == 'test4'
assert user_data[54321] == {}
chat_data = pickle_persistence.get_chat_data() chat_data = pickle_persistence.get_chat_data()
assert isinstance(chat_data, defaultdict) assert isinstance(chat_data, dict)
assert chat_data[-12345]['test1'] == 'test2' assert chat_data[-12345]['test1'] == 'test2'
assert chat_data[-67890][3] == 'test4' assert chat_data[-67890][3] == 'test4'
assert chat_data[-54321] == {}
bot_data = pickle_persistence.get_bot_data() bot_data = pickle_persistence.get_bot_data()
assert isinstance(bot_data, dict) assert isinstance(bot_data, dict)
@ -1310,16 +1293,14 @@ class TestPicklePersistence:
self, pickle_persistence, pickle_files_wo_callback_data self, pickle_persistence, pickle_files_wo_callback_data
): ):
user_data = pickle_persistence.get_user_data() user_data = pickle_persistence.get_user_data()
assert isinstance(user_data, defaultdict) assert isinstance(user_data, dict)
assert user_data[12345]['test1'] == 'test2' assert user_data[12345]['test1'] == 'test2'
assert user_data[67890][3] == 'test4' assert user_data[67890][3] == 'test4'
assert user_data[54321] == {}
chat_data = pickle_persistence.get_chat_data() chat_data = pickle_persistence.get_chat_data()
assert isinstance(chat_data, defaultdict) assert isinstance(chat_data, dict)
assert chat_data[-12345]['test1'] == 'test2' assert chat_data[-12345]['test1'] == 'test2'
assert chat_data[-67890][3] == 'test4' assert chat_data[-67890][3] == 'test4'
assert chat_data[-54321] == {}
bot_data = pickle_persistence.get_bot_data() bot_data = pickle_persistence.get_bot_data()
assert isinstance(bot_data, dict) assert isinstance(bot_data, dict)
@ -1353,7 +1334,7 @@ class TestPicklePersistence:
pickle_persistence.update_user_data(12345, user_data[12345]) pickle_persistence.update_user_data(12345, user_data[12345])
assert pickle_persistence.user_data == user_data assert pickle_persistence.user_data == user_data
with Path('pickletest_user_data').open('rb') as f: with Path('pickletest_user_data').open('rb') as f:
user_data_test = defaultdict(dict, pickle.load(f)) user_data_test = dict(pickle.load(f))
assert user_data_test == user_data assert user_data_test == user_data
pickle_persistence.drop_user_data(67890) pickle_persistence.drop_user_data(67890)
assert 67890 not in pickle_persistence.get_user_data() assert 67890 not in pickle_persistence.get_user_data()
@ -1367,7 +1348,7 @@ class TestPicklePersistence:
pickle_persistence.update_chat_data(-12345, chat_data[-12345]) pickle_persistence.update_chat_data(-12345, chat_data[-12345])
assert pickle_persistence.chat_data == chat_data assert pickle_persistence.chat_data == chat_data
with Path('pickletest_chat_data').open('rb') as f: with Path('pickletest_chat_data').open('rb') as f:
chat_data_test = defaultdict(dict, pickle.load(f)) chat_data_test = dict(pickle.load(f))
assert chat_data_test == chat_data assert chat_data_test == chat_data
pickle_persistence.drop_chat_data(-67890) pickle_persistence.drop_chat_data(-67890)
assert -67890 not in pickle_persistence.get_chat_data() assert -67890 not in pickle_persistence.get_chat_data()
@ -1403,7 +1384,7 @@ class TestPicklePersistence:
assert pickle_persistence.conversations['name1'] == conversation1 assert pickle_persistence.conversations['name1'] == conversation1
assert pickle_persistence.get_conversations('name1') == conversation1 assert pickle_persistence.get_conversations('name1') == conversation1
with Path('pickletest_conversations').open('rb') as f: with Path('pickletest_conversations').open('rb') as f:
conversations_test = defaultdict(dict, pickle.load(f)) conversations_test = dict(pickle.load(f))
assert conversations_test['name1'] == conversation1 assert conversations_test['name1'] == conversation1
pickle_persistence.conversations = None pickle_persistence.conversations = None
@ -1423,7 +1404,7 @@ class TestPicklePersistence:
pickle_persistence.update_user_data(12345, user_data[12345]) pickle_persistence.update_user_data(12345, user_data[12345])
assert pickle_persistence.user_data == user_data assert pickle_persistence.user_data == user_data
with Path('pickletest').open('rb') as f: with Path('pickletest').open('rb') as f:
user_data_test = defaultdict(dict, pickle.load(f)['user_data']) user_data_test = dict(pickle.load(f))['user_data']
assert user_data_test == user_data assert user_data_test == user_data
pickle_persistence.drop_user_data(67890) pickle_persistence.drop_user_data(67890)
assert 67890 not in pickle_persistence.get_user_data() assert 67890 not in pickle_persistence.get_user_data()
@ -1437,7 +1418,7 @@ class TestPicklePersistence:
pickle_persistence.update_chat_data(-12345, chat_data[-12345]) pickle_persistence.update_chat_data(-12345, chat_data[-12345])
assert pickle_persistence.chat_data == chat_data assert pickle_persistence.chat_data == chat_data
with Path('pickletest').open('rb') as f: with Path('pickletest').open('rb') as f:
chat_data_test = defaultdict(dict, pickle.load(f)['chat_data']) chat_data_test = dict(pickle.load(f))['chat_data']
assert chat_data_test == chat_data assert chat_data_test == chat_data
pickle_persistence.drop_chat_data(-67890) pickle_persistence.drop_chat_data(-67890)
assert -67890 not in pickle_persistence.get_chat_data() assert -67890 not in pickle_persistence.get_chat_data()
@ -1473,7 +1454,7 @@ class TestPicklePersistence:
assert pickle_persistence.conversations['name1'] == conversation1 assert pickle_persistence.conversations['name1'] == conversation1
assert pickle_persistence.get_conversations('name1') == conversation1 assert pickle_persistence.get_conversations('name1') == conversation1
with Path('pickletest').open('rb') as f: with Path('pickletest').open('rb') as f:
conversations_test = defaultdict(dict, pickle.load(f)['conversations']) conversations_test = dict(pickle.load(f))['conversations']
assert conversations_test['name1'] == conversation1 assert conversations_test['name1'] == conversation1
pickle_persistence.conversations = None pickle_persistence.conversations = None
@ -1502,6 +1483,7 @@ class TestPicklePersistence:
pickle_persistence.on_flush = True pickle_persistence.on_flush = True
user_data = pickle_persistence.get_user_data() user_data = pickle_persistence.get_user_data()
user_data[54321] = {}
user_data[54321]['test9'] = 'test 10' user_data[54321]['test9'] = 'test 10'
assert not pickle_persistence.user_data == user_data assert not pickle_persistence.user_data == user_data
@ -1512,10 +1494,11 @@ class TestPicklePersistence:
assert pickle_persistence.user_data == user_data assert pickle_persistence.user_data == user_data
with Path('pickletest_user_data').open('rb') as f: with Path('pickletest_user_data').open('rb') as f:
user_data_test = defaultdict(dict, pickle.load(f)) user_data_test = dict(pickle.load(f))
assert not user_data_test == user_data assert not user_data_test == user_data
chat_data = pickle_persistence.get_chat_data() chat_data = pickle_persistence.get_chat_data()
chat_data[54321] = {}
chat_data[54321]['test9'] = 'test 10' chat_data[54321]['test9'] = 'test 10'
assert not pickle_persistence.chat_data == chat_data assert not pickle_persistence.chat_data == chat_data
@ -1526,7 +1509,7 @@ class TestPicklePersistence:
assert pickle_persistence.user_data == user_data assert pickle_persistence.user_data == user_data
with Path('pickletest_chat_data').open('rb') as f: with Path('pickletest_chat_data').open('rb') as f:
chat_data_test = defaultdict(dict, pickle.load(f)) chat_data_test = dict(pickle.load(f))
assert not chat_data_test == chat_data assert not chat_data_test == chat_data
bot_data = pickle_persistence.get_bot_data() bot_data = pickle_persistence.get_bot_data()
@ -1559,16 +1542,16 @@ class TestPicklePersistence:
assert pickle_persistence.conversations['name1'] == conversation1 assert pickle_persistence.conversations['name1'] == conversation1
with Path('pickletest_conversations').open('rb') as f: with Path('pickletest_conversations').open('rb') as f:
conversations_test = defaultdict(dict, pickle.load(f)) conversations_test = dict(pickle.load(f))
assert not conversations_test['name1'] == conversation1 assert not conversations_test['name1'] == conversation1
pickle_persistence.flush() pickle_persistence.flush()
with Path('pickletest_user_data').open('rb') as f: with Path('pickletest_user_data').open('rb') as f:
user_data_test = defaultdict(dict, pickle.load(f)) user_data_test = dict(pickle.load(f))
assert user_data_test == user_data assert user_data_test == user_data
with Path('pickletest_chat_data').open('rb') as f: with Path('pickletest_chat_data').open('rb') as f:
chat_data_test = defaultdict(dict, pickle.load(f)) chat_data_test = dict(pickle.load(f))
assert chat_data_test == chat_data assert chat_data_test == chat_data
with Path('pickletest_bot_data').open('rb') as f: with Path('pickletest_bot_data').open('rb') as f:
@ -1576,7 +1559,7 @@ class TestPicklePersistence:
assert bot_data_test == bot_data assert bot_data_test == bot_data
with Path('pickletest_conversations').open('rb') as f: with Path('pickletest_conversations').open('rb') as f:
conversations_test = defaultdict(dict, pickle.load(f)) conversations_test = dict(pickle.load(f))
assert conversations_test['name1'] == conversation1 assert conversations_test['name1'] == conversation1
def test_save_on_flush_single_files(self, pickle_persistence, good_pickle_files): def test_save_on_flush_single_files(self, pickle_persistence, good_pickle_files):
@ -1587,21 +1570,23 @@ class TestPicklePersistence:
pickle_persistence.single_file = True pickle_persistence.single_file = True
user_data = pickle_persistence.get_user_data() user_data = pickle_persistence.get_user_data()
user_data[54321] = {}
user_data[54321]['test9'] = 'test 10' user_data[54321]['test9'] = 'test 10'
assert not pickle_persistence.user_data == user_data assert not pickle_persistence.user_data == user_data
pickle_persistence.update_user_data(54321, user_data[54321]) pickle_persistence.update_user_data(54321, user_data[54321])
assert pickle_persistence.user_data == user_data assert pickle_persistence.user_data == user_data
with Path('pickletest').open('rb') as f: with Path('pickletest').open('rb') as f:
user_data_test = defaultdict(dict, pickle.load(f)['user_data']) user_data_test = dict(pickle.load(f))['user_data']
assert not user_data_test == user_data assert not user_data_test == user_data
chat_data = pickle_persistence.get_chat_data() chat_data = pickle_persistence.get_chat_data()
chat_data[54321] = {}
chat_data[54321]['test9'] = 'test 10' chat_data[54321]['test9'] = 'test 10'
assert not pickle_persistence.chat_data == chat_data assert not pickle_persistence.chat_data == chat_data
pickle_persistence.update_chat_data(54321, chat_data[54321]) pickle_persistence.update_chat_data(54321, chat_data[54321])
assert pickle_persistence.chat_data == chat_data assert pickle_persistence.chat_data == chat_data
with Path('pickletest').open('rb') as f: with Path('pickletest').open('rb') as f:
chat_data_test = defaultdict(dict, pickle.load(f)['chat_data']) chat_data_test = dict(pickle.load(f))['chat_data']
assert not chat_data_test == chat_data assert not chat_data_test == chat_data
bot_data = pickle_persistence.get_bot_data() bot_data = pickle_persistence.get_bot_data()
@ -1628,16 +1613,16 @@ class TestPicklePersistence:
pickle_persistence.update_conversation('name1', (123, 123), 5) pickle_persistence.update_conversation('name1', (123, 123), 5)
assert pickle_persistence.conversations['name1'] == conversation1 assert pickle_persistence.conversations['name1'] == conversation1
with Path('pickletest').open('rb') as f: with Path('pickletest').open('rb') as f:
conversations_test = defaultdict(dict, pickle.load(f)['conversations']) conversations_test = dict(pickle.load(f))['conversations']
assert not conversations_test['name1'] == conversation1 assert not conversations_test['name1'] == conversation1
pickle_persistence.flush() pickle_persistence.flush()
with Path('pickletest').open('rb') as f: with Path('pickletest').open('rb') as f:
user_data_test = defaultdict(dict, pickle.load(f)['user_data']) user_data_test = dict(pickle.load(f))['user_data']
assert user_data_test == user_data assert user_data_test == user_data
with Path('pickletest').open('rb') as f: with Path('pickletest').open('rb') as f:
chat_data_test = defaultdict(dict, pickle.load(f)['chat_data']) chat_data_test = dict(pickle.load(f))['chat_data']
assert chat_data_test == chat_data assert chat_data_test == chat_data
with Path('pickletest').open('rb') as f: with Path('pickletest').open('rb') as f:
@ -1645,7 +1630,7 @@ class TestPicklePersistence:
assert bot_data_test == bot_data assert bot_data_test == bot_data
with Path('pickletest').open('rb') as f: with Path('pickletest').open('rb') as f:
conversations_test = defaultdict(dict, pickle.load(f)['conversations']) conversations_test = dict(pickle.load(f))['conversations']
assert conversations_test['name1'] == conversation1 assert conversations_test['name1'] == conversation1
def test_with_handler(self, bot, update, bot_data, pickle_persistence, good_pickle_files): def test_with_handler(self, bot, update, bot_data, pickle_persistence, good_pickle_files):
@ -1925,10 +1910,6 @@ class TestPicklePersistence:
cc = ContextTypes(user_data=ud, chat_data=cd, bot_data=bd) cc = ContextTypes(user_data=ud, chat_data=cd, bot_data=bd)
persistence = PicklePersistence('pickletest', single_file=singlefile, context_types=cc) persistence = PicklePersistence('pickletest', single_file=singlefile, context_types=cc)
assert isinstance(persistence.get_user_data()[1], ud)
assert persistence.get_user_data()[1] == 0
assert isinstance(persistence.get_chat_data()[1], cd)
assert persistence.get_chat_data()[1] == 0
assert isinstance(persistence.get_bot_data(), bd) assert isinstance(persistence.get_bot_data(), bd)
assert persistence.get_bot_data() == 0 assert persistence.get_bot_data() == 0
@ -1936,8 +1917,8 @@ class TestPicklePersistence:
persistence.chat_data = None persistence.chat_data = None
persistence.drop_user_data(123) persistence.drop_user_data(123)
persistence.drop_chat_data(123) persistence.drop_chat_data(123)
assert isinstance(persistence.get_user_data(), defaultdict) assert isinstance(persistence.get_user_data(), dict)
assert isinstance(persistence.get_chat_data(), defaultdict) assert isinstance(persistence.get_chat_data(), dict)
persistence.user_data = None persistence.user_data = None
persistence.chat_data = None persistence.chat_data = None
persistence.update_user_data(1, ud(1)) persistence.update_user_data(1, ud(1))
@ -1993,8 +1974,8 @@ class TestDictPersistence:
def test_no_json_given(self): def test_no_json_given(self):
dict_persistence = DictPersistence() dict_persistence = DictPersistence()
assert dict_persistence.get_user_data() == defaultdict(dict) assert dict_persistence.get_user_data() == {}
assert dict_persistence.get_chat_data() == defaultdict(dict) assert dict_persistence.get_chat_data() == {}
assert dict_persistence.get_bot_data() == {} assert dict_persistence.get_bot_data() == {}
assert dict_persistence.get_callback_data() is None assert dict_persistence.get_callback_data() is None
assert dict_persistence.get_conversations('noname') == {} assert dict_persistence.get_conversations('noname') == {}
@ -2055,16 +2036,14 @@ class TestDictPersistence:
callback_data_json=callback_data_json, callback_data_json=callback_data_json,
) )
user_data = dict_persistence.get_user_data() user_data = dict_persistence.get_user_data()
assert isinstance(user_data, defaultdict) assert isinstance(user_data, dict)
assert user_data[12345]['test1'] == 'test2' assert user_data[12345]['test1'] == 'test2'
assert user_data[67890][3] == 'test4' assert user_data[67890][3] == 'test4'
assert user_data[54321] == {}
chat_data = dict_persistence.get_chat_data() chat_data = dict_persistence.get_chat_data()
assert isinstance(chat_data, defaultdict) assert isinstance(chat_data, dict)
assert chat_data[-12345]['test1'] == 'test2' assert chat_data[-12345]['test1'] == 'test2'
assert chat_data[-67890][3] == 'test4' assert chat_data[-67890][3] == 'test4'
assert chat_data[-54321] == {}
bot_data = dict_persistence.get_bot_data() bot_data = dict_persistence.get_bot_data()
assert isinstance(bot_data, dict) assert isinstance(bot_data, dict)
@ -2169,7 +2148,7 @@ class TestDictPersistence:
assert 67890 not in dict_persistence.user_data assert 67890 not in dict_persistence.user_data
dict_persistence._user_data = None dict_persistence._user_data = None
dict_persistence.drop_user_data(123) dict_persistence.drop_user_data(123)
assert isinstance(dict_persistence.get_user_data(), defaultdict) assert isinstance(dict_persistence.get_user_data(), dict)
chat_data = dict_persistence.get_chat_data() chat_data = dict_persistence.get_chat_data()
chat_data[-12345]['test3']['test4'] = 'test6' chat_data[-12345]['test3']['test4'] = 'test6'
@ -2186,7 +2165,7 @@ class TestDictPersistence:
assert -67890 not in dict_persistence.chat_data assert -67890 not in dict_persistence.chat_data
dict_persistence._chat_data = None dict_persistence._chat_data = None
dict_persistence.drop_chat_data(123) dict_persistence.drop_chat_data(123)
assert isinstance(dict_persistence.get_chat_data(), defaultdict) assert isinstance(dict_persistence.get_chat_data(), dict)
bot_data = dict_persistence.get_bot_data() bot_data = dict_persistence.get_bot_data()
bot_data['test3']['test4'] = 'test6' bot_data['test3']['test4'] = 'test6'