Merge pull request #68 from rhutkovich/feature/async-django-impl

Made methods async, same as the most recent BasePersistence class

Author: GamePad64

python_telegram_bot_django_persistence/persistence.py

@@ -2,105 +2,104 @@
 from collections import defaultdict
 from typing import DefaultDict, Optional, Tuple, cast
 
-from telegram.ext import BasePersistence
-from telegram.ext.utils.types import BD, CD, UD, CDCData, ConversationDict
-
-from .models import BotData, CallbackData, ChatData, ConversationData, UserData
+from models import BotData, CallbackData, ChatData, ConversationData, UserData
+from telegram.ext import BasePersistence, PersistenceInput
+from telegram.ext._utils.types import BD, CD, UD, CDCData, ConversationDict
 
 
 class DjangoPersistence(BasePersistence[UD, CD, BD]):
     def __init__(
         self,
         namespace: str = "",
-        store_user_data: bool = True,
-        store_chat_data: bool = True,
-        store_bot_data: bool = True,
-        store_callback_data: bool = False,
+        store_data: Optional[PersistenceInput] = None,
+        update_interval: float = 60,
     ):
-        super().__init__(
-            store_user_data=store_user_data,
-            store_chat_data=store_chat_data,
-            store_bot_data=store_bot_data,
-            store_callback_data=store_callback_data,
-        )
+        super().__init__(store_data=store_data, update_interval=update_interval)
         self._namespace = namespace
 
-    def get_bot_data(self) -> BD:
+    async def get_bot_data(self) -> BD:
         try:
-            return BotData.objects.get(namespace=self._namespace).data
+            return (await BotData.objects.aget(namespace=self._namespace)).data
         except BotData.DoesNotExist:
             return {}
 
-    def update_bot_data(self, data: BD) -> None:
-        BotData.objects.update_or_create(namespace=self._namespace, defaults={"data": data})
+    async def update_bot_data(self, data: BD) -> None:
+        await BotData.objects.aupdate_or_create(namespace=self._namespace, defaults={"data": data})
 
-    def refresh_bot_data(self, bot_data: BD) -> None:
+    async def refresh_bot_data(self, bot_data: BD) -> None:
         if isinstance(bot_data, dict):
             orig_keys = set(bot_data.keys())
-            bot_data.update(self.get_bot_data())
+            bot_data.update(await self.get_bot_data())
             for key in orig_keys - set(bot_data.keys()):
                 bot_data.pop(key)
 
-    def get_chat_data(self) -> DefaultDict[int, CD]:
-        return defaultdict(
-            dict, {data.chat_id: data.data for data in ChatData.objects.filter(namespace=self._namespace)}
-        )
+    async def get_chat_data(self) -> DefaultDict[int, CD]:
+        chat_data = {}
+        async for data in ChatData.objects.filter(namespace=self._namespace):
+            chat_data[data.chat_id] = data.data
+        return defaultdict(dict, chat_data)
 
-    def update_chat_data(self, chat_id: int, data: CD) -> None:
-        ChatData.objects.update_or_create(namespace=self._namespace, chat_id=chat_id, defaults={"data": data})
+    async def update_chat_data(self, chat_id: int, data: CD) -> None:
+        await ChatData.objects.aupdate_or_create(namespace=self._namespace, chat_id=chat_id, defaults={"data": data})
 
-    def refresh_chat_data(self, chat_id: int, chat_data: CD) -> None:
+    async def refresh_chat_data(self, chat_id: int, chat_data: CD) -> None:
         try:
             if isinstance(chat_data, dict):
                 orig_keys = set(chat_data.keys())
-                chat_data.update(ChatData.objects.get(namespace=self._namespace, chat_id=chat_id).data)
+                chat_data.update((await ChatData.objects.aget(namespace=self._namespace, chat_id=chat_id)).data)
                 for key in orig_keys - set(chat_data.keys()):
                     chat_data.pop(key)
         except ChatData.DoesNotExist:
             pass
 
-    def get_user_data(self) -> DefaultDict[int, UD]:
-        return defaultdict(
-            dict, {data.user_id: data.data for data in UserData.objects.filter(namespace=self._namespace)}
-        )
+    async def get_user_data(self) -> DefaultDict[int, UD]:
+        user_data = {}
+        async for data in UserData.objects.filter(namespace=self._namespace):
+            user_data[data.user_id] = data.data
+        return defaultdict(dict, user_data)
 
-    def update_user_data(self, user_id: int, data: UD) -> None:
-        UserData.objects.update_or_create(namespace=self._namespace, user_id=user_id, defaults={"data": data})
+    async def update_user_data(self, user_id: int, data: UD) -> None:
+        await UserData.objects.aupdate_or_create(namespace=self._namespace, user_id=user_id, defaults={"data": data})
 
-    def refresh_user_data(self, user_id: int, user_data: UD) -> None:
+    async def refresh_user_data(self, user_id: int, user_data: UD) -> None:
         try:
             if isinstance(user_data, dict):
                 orig_keys = set(user_data.keys())
-                user_data.update(UserData.objects.get(namespace=self._namespace, user_id=user_id).data)
+                user_data.update((await UserData.objects.aget(namespace=self._namespace, user_id=user_id)).data)
                 for key in orig_keys - set(user_data.keys()):
                     user_data.pop(key)
         except UserData.DoesNotExist:
             pass
 
-    def get_callback_data(self) -> Optional[CDCData]:
+    async def get_callback_data(self) -> Optional[CDCData]:
         try:
-            cdcdata_json = CallbackData.objects.get(namespace=self._namespace).data
+            cdcdata_json = (await CallbackData.objects.aget(namespace=self._namespace)).data
             # Before asking me wtf is this, just check DictPersistence
             return cast(CDCData, ([(one, float(two), three) for one, two, three in cdcdata_json[0]], cdcdata_json[1]))
         except CallbackData.DoesNotExist:
             return None
 
-    def update_callback_data(self, data: CDCData) -> None:
-        CallbackData.objects.update_or_create(namespace=self._namespace, defaults={"data": data})
+    async def update_callback_data(self, data: CDCData) -> None:
+        await CallbackData.objects.aupdate_or_create(namespace=self._namespace, defaults=data)
 
-    def get_conversations(self, name: str) -> ConversationDict:
-        return {
-            tuple(json.loads(data.key)): data.state
-            for data in ConversationData.objects.filter(namespace=self._namespace, name=name)
-        }
+    async def get_conversations(self, name: str) -> ConversationDict:
+        result = {}
+        async for data in ConversationData.objects.filter(namespace=self._namespace, name=name):
+            result[tuple(json.loads(data.key))] = data.state
 
-    def update_conversation(self, name: str, key: Tuple[int, ...], new_state: Optional[object]) -> None:
-        ConversationData.objects.update_or_create(
+    async def update_conversation(self, name: str, key: Tuple[int, ...], new_state: Optional[object]) -> None:
+        await ConversationData.objects.aupdate_or_create(
             namespace=self._namespace,
             name=name,
             key=json.dumps(key, sort_keys=True),
             defaults={"state": new_state},
         )
 
-    def flush(self) -> None:
+    async def flush(self) -> None:
         pass
+
+    async def drop_chat_data(self, chat_id: int) -> None:
+        await ChatData.objects.filter(namespace=self._namespace, chat_id=chat_id).adelete()
+
+    async def drop_user_data(self, user_id: int) -> None:
+        await UserData.objects.filter(namespace=self._namespace, user_id=user_id).adelete()