add room & contact cache

Author: wj-Mcat

src/wechaty/user/contact.py

@@ -21,6 +21,7 @@
 from __future__ import annotations
 
 import types
+import os
 import asyncio
 import dataclasses
 import json
@@ -47,8 +48,9 @@
 
 # from wechaty.utils import type_check
 from wechaty.exceptions import WechatyPayloadError, WechatyOperationError
-from wechaty.config import PARALLEL_TASK_NUM
+from wechaty.config import PARALLEL_TASK_NUM, config
 from wechaty.utils.async_helper import gather_with_concurrency
+from wechaty.utils.data_util import save_pickle_data, load_pickle_data
 
 from ..accessory import Accessory
 
@@ -64,7 +66,6 @@
 
 log = get_logger('Contact')
 
-
 # pylint:disable=R0904
 class Contact(Accessory[ContactPayload], AsyncIOEventEmitter):
     """
@@ -250,18 +251,39 @@ async def find_all(cls: Type[Contact],
         """
         log.info('find_all() <%s, %s>', cls, query)
 
+        # 0. load from local cache file
+        if config.cache_contacts and os.path.exists(config.cache_contact_path):
+            contact_payloads = load_pickle_data(config.cache_contact_path)
+            assert isinstance(contact_payloads, list)
+            assert isinstance(contact_payloads[0], ContactPayload)
+
+            contact_payloads: List[ContactPayload] = contact_payloads
+            contacts = []
+            for contact_payload in contact_payloads:
+                contact = cls.load(contact_payload.id)
+
+                # type: ignore
+                contact._payload = contact_payload
+                contacts.append(contact)
+            return contacts
+
         # 1. load contacts with concurrent tasks
         contact_ids: List[str] = await cls.get_puppet().contact_list()
 
         contacts: List[Contact] = [cls.load(contact_id) for contact_id in contact_ids]
         tasks: List[Task] = [asyncio.create_task(contact.ready()) for contact in contacts]
         await gather_with_concurrency(PARALLEL_TASK_NUM, tasks)
 
+        if config.cache_contacts:
+            contact_payloads = [contact.payload for contact in contacts if contact.payload is not None]
+            save_pickle_data(contact_payloads, config.cache_contact_path)
+
         # 2. filter contacts
         if not query:
             return contacts
 
         contacts = cls._filter_contacts(contacts, query)
+
         return contacts
 
     async def ready(self, force_sync: bool = False) -> None:

src/wechaty/user/room.py

@@ -19,6 +19,7 @@
 limitations under the License.
 """
 from __future__ import annotations
+import os
 import types
 import asyncio
 import dataclasses
@@ -38,6 +39,7 @@
 import json
 from pyee import AsyncIOEventEmitter
 from wechaty.exceptions import WechatyOperationError, WechatyPayloadError
+from wechaty.config import config
 from wechaty_puppet import (
     FileBox,
     RoomQueryFilter,
@@ -50,6 +52,7 @@
 # from wechaty.utils import type_check
 from ..accessory import Accessory
 from ..config import AT_SEPARATOR, PARALLEL_TASK_NUM
+from wechaty.utils.data_util import save_pickle_data, load_pickle_data
 from wechaty.utils.async_helper import gather_with_concurrency
 
 if TYPE_CHECKING:
@@ -222,11 +225,31 @@ async def find_all(cls,
         """
         log.info('Room find_all <%s>', query)
 
-        # 1. load rooms with concurrent tasks
-        room_ids = await cls.get_puppet().room_search()
-        rooms: List[Room] = [cls.load(room_id) for room_id in room_ids]
-        tasks: List[Task] = [asyncio.create_task(room.ready()) for room in rooms]
-        await gather_with_concurrency(PARALLEL_TASK_NUM, tasks)
+        # 0. load from local cache file
+        if config.cache_rooms and os.path.exists(config.cache_room_path):
+            room_payloads = load_pickle_data(config.cache_room_path)
+            assert isinstance(room_payloads, list)
+            assert isinstance(room_payloads[0], RoomPayload)
+
+            room_payloads: List[RoomPayload] = room_payloads
+            rooms = []
+            for room_payload in room_payloads:
+                room = cls.load(room_payload.id)
+
+                # type: ignore
+                room._payload = room_payload
+                rooms.append(room)
+        else:
+
+            # 1. load rooms with concurrent tasks
+            room_ids = await cls.get_puppet().room_search()
+            rooms: List[Room] = [cls.load(room_id) for room_id in room_ids]
+            tasks: List[Task] = [asyncio.create_task(room.ready()) for room in rooms]
+            await gather_with_concurrency(PARALLEL_TASK_NUM, tasks)
+        
+        if config.cache_rooms:
+            room_payloads = [room.payload for room in rooms if room.payload is not None]
+            save_pickle_data(room_payloads, config.cache_contact_path)
 
         # 2. filter the rooms
         if not query:

src/wechaty/utils/data_util.py

@@ -26,7 +26,31 @@
     Any,
 )
 from collections import UserDict
-from wechaty.config import config
+import pickle
+
+
+def save_pickle_data(obj: object, path: str):
+    """save pickle data
+
+    Args:
+        obj (object): the pickled data
+        path (str): the path of pickle data
+    """
+    with open(path, 'wb') as f:
+        pickle.dump(obj, f)
+
+def load_pickle_data(path: str) -> object:
+    """load pickle data from path
+
+    Args:
+        path (str): the path of pickle data
+
+    Returns:
+        object: the final data
+    """
+    with open(path, 'rb') as f:
+        data = pickle.load(f)
+    return data
 
 
 class WechatySetting(UserDict):