From 7f532205249692f4143648057193fc360f97bdd6 Mon Sep 17 00:00:00 2001 From: hikariatama Date: Tue, 9 Aug 2022 22:05:35 +0000 Subject: [PATCH] Add `client.get_fulluser` with cache --- CHANGELOG.md | 1 + hikka/main.py | 8 +++- hikka/tl_cache.py | 96 +++++++++++++++++++++++++++++++++++++++++++++-- 3 files changed, 101 insertions(+), 4 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 8ed7e09..8266ba2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -16,6 +16,7 @@ - Add possible cause of error in logs (module and method) - Add `client.get_perms_cached` to cache native `client.get_permissions` - Add `client.get_fullchannel` with cache +- Add `client.get_fulluser` with cache - Add `exp` and `force` params to `client.get_perms_cached` and `client.get_entity` - Add `exp` cached values check in `client.get_perms_cached` and `client.get_entity` - Change errors format in web to more human-readable diff --git a/hikka/main.py b/hikka/main.py index 3ee19e7..97df08c 100755 --- a/hikka/main.py +++ b/hikka/main.py @@ -56,7 +56,12 @@ from . import database, loader, utils, heroku from .dispatcher import CommandDispatcher from .translations import Translator from .version import __version__ -from .tl_cache import install_entity_caching, install_perms_caching, install_fullchannel_caching +from .tl_cache import ( + install_entity_caching, + install_perms_caching, + install_fullchannel_caching, + install_fulluser_caching, +) try: from .web import core @@ -526,6 +531,7 @@ class Hikka: install_entity_caching(client) install_perms_caching(client) install_fullchannel_caching(client) + install_fulluser_caching(client) self.clients += [client] except sqlite3.OperationalError: diff --git a/hikka/tl_cache.py b/hikka/tl_cache.py index cb5722e..9315d56 100644 --- a/hikka/tl_cache.py +++ b/hikka/tl_cache.py @@ -11,11 +11,11 @@ import time import asyncio import logging from typing import Optional, Union -from xml.dom.minidom import Entity from telethon.hints import EntityLike from telethon import TelegramClient from telethon.tl.functions.channels import GetFullChannelRequest -from telethon.tl.types import ChannelFull +from telethon.tl.functions.users import GetFullUserRequest +from telethon.tl.types import ChannelFull, UserFull logger = logging.getLogger(__name__) @@ -116,6 +116,33 @@ class CacheRecordFullChannel: ) + +class CacheRecordFullUser: + def __init__(self, user_id: int, full_user: UserFull, exp: int): + self.user_id = user_id + self.full_user = full_user + self._exp = round(time.time() + exp) + self.ts = time.time() + + def expired(self): + return self._exp < time.time() + + def __eq__(self, record: "CacheRecordFullUser"): + return hash(record) == hash(self) + + def __hash__(self): + return hash((self._hashable_entity, self._hashable_user)) + + def __str__(self): + return f"CacheRecordFullUser of {self.user_id}" + + def __repr__(self): + return ( + f"CacheRecordFullUser(channel_id={self.user_id}(...)," + f" exp={self._exp})" + ) + + def install_entity_caching(client: TelegramClient): client._hikka_entity_cache = {} @@ -337,7 +364,7 @@ def install_fullchannel_caching(client: TelegramClient): :param entity: Channel to fetch ChannelFull of :param exp: Expiration time of the cache record and maximum time of already cached record :param force: Whether to force refresh the cache (make API request) - :return: :obj:`FullChannel` + :return: :obj:`ChannelFull` """ if not hashable(entity): try: @@ -385,3 +412,66 @@ def install_fullchannel_caching(client: TelegramClient): client.get_fullchannel = get_fullchannel asyncio.ensure_future(cleaner(client)) logger.debug("Monkeypatched client with fullchannel cacher") + + +def install_fulluser_caching(client: TelegramClient): + client._hikka_fulluser_cache = {} + + async def get_fulluser( + entity: EntityLike, + exp: Optional[int] = 300, + force: Optional[bool] = False, + ) -> ChannelFull: + """ + Gets the FullUserRequest and cache it + :param entity: User to fetch UserFull of + :param exp: Expiration time of the cache record and maximum time of already cached record + :param force: Whether to force refresh the cache (make API request) + :return: :obj:`UserFull` + """ + if not hashable(entity): + try: + hashable_entity = next( + getattr(entity, attr) + for attr in {"user_id", "chat_id", "id"} + if getattr(entity, attr, None) + ) + except StopIteration: + logger.debug( + f"Can't parse hashable from {entity=}, using legacy fulluser request" + ) + return await client(GetFullUserRequest(entity)) + else: + hashable_entity = entity + + if str(hashable_entity).isdigit() and int(hashable_entity) < 0: + hashable_entity = int(str(hashable_entity)[4:]) + + if ( + not force + and client._hikka_fulluser_cache.get(hashable_entity) + and not client._hikka_fulluser_cache[hashable_entity].expired() + and client._hikka_fulluser_cache[hashable_entity].ts + exp > time.time() + ): + return client._hikka_fulluser_cache[hashable_entity].full_channel + + result = await client(GetFullUserRequest(entity)) + client._hikka_fulluser_cache[hashable_entity] = CacheRecordFullUser( + hashable_entity, + result, + exp, + ) + return result + + async def cleaner(client: TelegramClient): + while True: + for user_id, record in client._hikka_fulluser_cache.copy().items(): + if record.expired(): + del client._hikka_fulluser_cache[user_id] + logger.debug(f"Cleaned outdated fulluser cache {user_id=}") + + await asyncio.sleep(3) + + client.get_fulluser = get_fulluser + asyncio.ensure_future(cleaner(client)) + logger.debug("Monkeypatched client with fulluser cacher")