diff options
| author | Nakidai <nakidai@disroot.org> | 2024-11-07 03:02:00 +0300 |
|---|---|---|
| committer | Nakidai <nakidai@disroot.org> | 2024-11-07 03:02:00 +0300 |
| commit | 4f99e6ada3ddde9c8b1a79129db7653e5862c985 (patch) | |
| tree | 0c8e5445c23cf23c292b1a4bac8e063d257f0da2 | |
| parent | dda6671fdbbbd6b3a1f2529dbc60726b67cd4fb2 (diff) | |
| download | petthecord-4f99e6ada3ddde9c8b1a79129db7653e5862c985.tar.gz petthecord-4f99e6ada3ddde9c8b1a79129db7653e5862c985.zip | |
Refactor cache
- Now runner only gets petter as the arguemnt, not petter's arguemnts - Petters has its own directory
| -rw-r--r-- | src/petthecord/__init__.py | 3 | ||||
| -rw-r--r-- | src/petthecord/cache.py | 121 | ||||
| -rw-r--r-- | src/petthecord/main.py | 15 | ||||
| -rw-r--r-- | src/petthecord/petter/__init__.py | 6 | ||||
| -rw-r--r-- | src/petthecord/petter/cachedpetter.py | 97 | ||||
| -rw-r--r-- | src/petthecord/petter/exceptions.py | 10 | ||||
| -rw-r--r-- | src/petthecord/petter/petter.py | 45 | ||||
| -rw-r--r-- | src/petthecord/runner.py | 27 | ||||
| -rw-r--r-- | src/petthecord/server.py | 4 |
9 files changed, 180 insertions, 148 deletions
diff --git a/src/petthecord/__init__.py b/src/petthecord/__init__.py index 65a78e9..9c17cd3 100644 --- a/src/petthecord/__init__.py +++ b/src/petthecord/__init__.py @@ -1,8 +1,7 @@ from .bot import PetTheCordCog -from .cache import CachedPet from .main import main from .runner import PetTheCord from .server import Server -__all__ = ["CachedPet", "main", "PetTheCord", "PetTheCordCog", "Server"] +__all__ = ["main", "PetTheCord", "PetTheCordCog", "Server"] diff --git a/src/petthecord/cache.py b/src/petthecord/cache.py deleted file mode 100644 index aa5618d..0000000 --- a/src/petthecord/cache.py +++ /dev/null @@ -1,121 +0,0 @@ -from asyncio import sleep -from json import load, dump -from io import BytesIO -from logging import getLogger -from os import PathLike, makedirs, remove, listdir -from os.path import getmtime -from pathlib import Path -from time import time -from typing import NoReturn - -import discord -from discord import Client -from petpetgif import petpet - - -class NotFound(Exception): - pass - - -class HTTPException(Exception): - pass - - -class CachedPet: - def __init__( - self, - - client: Client, - caching: bool = True, - path: str | PathLike = "/var/cache/petthecord", - lifetime: int = 86400, - gc_delay: int = 14400, - ) -> None: - self._client = client - self._caching = caching - self._path = Path(path).resolve() - self._lifetime = lifetime - self._gc_delay = gc_delay - - self._logger = getLogger(__name__) - - index_path = self._path / "index.json" - try: - if not index_path.exists(): - self._logger.warning("`index.json` doesn't exist, trying to create...") - if not self._path.exists(): - self._logger.warning("Cache folder doesnt exist, trying to create...") - makedirs(self._path, mode=0o755, exist_ok=True) - with open(index_path, "w") as f: - f.write("{}") - except OSError: - self._logger.error("Cannot create environment") - raise - - with open(index_path, "r") as f: - self._cache = load(f) - - async def petpet(self, uid: int) -> bytes: - try: - user = await self._client.fetch_user(uid) - except discord.NotFound: - raise NotFound - except discord.HTTPException: - raise HTTPException - - if user.avatar is None: - raise NotFound - - if self._caching: - avatar_path = self._path / f"{user.id}_{user.avatar.key}.gif" - if (path := self._cache.get(user.id)) != str(avatar_path): - self._logger.debug("Generating new gif for {user.id}") - if path: - try: - remove(path) - except OSError: - self._logger.warning("no {path} was found when replacing avatar") - self._cache[user.id] = str(avatar_path) - with open(self._path / "index.json", "w") as f: - dump(self._cache, f) - - if not avatar_path.exists(): - with open(avatar_path, "wb") as f: - image = await user.avatar.read() - petpet.make(BytesIO(image), f) - - avatar_path.touch() - - with open(avatar_path, "rb") as f: - return f.read() - else: - with BytesIO() as f: - image = await user.avatar.read() - petpet.make(BytesIO(image), f) - - f.seek(0) - - return f.read() - - async def gc_loop(self) -> NoReturn: - while True: - if self._caching: - self._logger.info("Starting new cache's gc iteration") - - for filename in listdir(self._path): - path = (self._path / filename) - if path.is_file() and filename != "index.json": - if (time() - getmtime(path) > self._lifetime): - self._logger.debug(f"Removing {filename}") - to_delete = filename.split('_')[0] - try: - del self._cache[to_delete] - except KeyError: - self._logger.warning(f"{to_delete} has been already removed from the index") - remove(path) - with open(self._path / "index.json", "w") as f: - dump(self._cache, f) - - self._logger.debug("Finished collecting old cache") - - await sleep(self._gc_delay) diff --git a/src/petthecord/main.py b/src/petthecord/main.py index 9177926..3ea56cb 100644 --- a/src/petthecord/main.py +++ b/src/petthecord/main.py @@ -2,6 +2,7 @@ from argparse import ArgumentParser from os import getenv from sys import argv, stderr +from .petter import CacheEnvironmentFail, CachedPetter, Petter from .runner import PetTheCord @@ -64,14 +65,20 @@ def main() -> None: ) args = parser.parse_args() + try: + petter = CachedPetter( + args.cache_dir, + args.cache_lifetime, + args.cache_gc_delay, + ) + except CacheEnvironmentFail: + petter = Petter() + bot = PetTheCord( args.host, args.port, args.origin, - not args.no_cache, - args.cache_dir, - args.cache_lifetime, - args.cache_gc_delay, + petter, args.shards, ) if (token := getenv("PETTHECORD_TOKEN")) is not None: diff --git a/src/petthecord/petter/__init__.py b/src/petthecord/petter/__init__.py new file mode 100644 index 0000000..3f3c6f1 --- /dev/null +++ b/src/petthecord/petter/__init__.py @@ -0,0 +1,6 @@ +from .cachedpetter import CachedPetter +from .exceptions import CacheEnvironmentFail, HTTPException, NotFound +from .petter import Petter + + +__all__ = ["CacheEnvironmentFail", "CachedPetter", "HTTPException", "NotFound", "Petter"] diff --git a/src/petthecord/petter/cachedpetter.py b/src/petthecord/petter/cachedpetter.py new file mode 100644 index 0000000..4e8b35d --- /dev/null +++ b/src/petthecord/petter/cachedpetter.py @@ -0,0 +1,97 @@ +from asyncio import sleep +from json import load, dump +from logging import getLogger +from os import PathLike, makedirs, remove, listdir +from os.path import getmtime +from pathlib import Path +from time import time +from typing import NoReturn + +from discord import User + +from .exceptions import CacheEnvironmentFail, NotFound +from .petter import Petter + + +class CachedPetter(Petter): + def __init__( + self, + + path: str | PathLike = "/var/cache/petthecord", + lifetime: int = 86400, + gc_delay: int = 14400, + ) -> None: + super().__init__() + + self._path = Path(path).resolve() + self._lifetime = lifetime + self._gc_delay = gc_delay + + self._logger = getLogger(__name__) + + index_path = self._path / "index.json" + try: + if not index_path.exists(): + self._logger.warning("`index.json` doesn't exist, trying to create...") + if not self._path.exists(): + self._logger.warning("Cache folder doesnt exist, trying to create...") + makedirs(self._path, mode=0o755, exist_ok=True) + with open(index_path, "w") as f: + f.write("{}") + except OSError: + self._logger.error("Cannot create environment") + raise CacheEnvironmentFail + + with open(index_path, "r") as f: + self._cache = load(f) + + async def petpet(self, user: int | User) -> bytes: + if isinstance(user, int): + _user = await self._get_user(user) + else: + _user = user + if not _user.avatar: + raise NotFound + + avatar_path = self._path / f"{_user.id}_{_user.avatar.key}.gif" + if (path := self._cache.get(_user.id)) != str(avatar_path): + self._logger.debug(f"Generating new cached gif for {_user.id}") + if path: + try: + remove(path) + except OSError: + self._logger.warning("no {path} was found when replacing avatar") + self._cache[_user.id] = str(avatar_path) + with open(self._path / "index.json", "w") as f: + dump(self._cache, f) + + if not avatar_path.exists(): + with open(avatar_path, "wb") as f: + f.write(await super().petpet(user)) + + avatar_path.touch() + + with open(avatar_path, "rb") as f: + return f.read() + + async def gc_loop(self) -> NoReturn: + while True: + self._logger.info("Starting new cache's gc iteration") + + for filename in listdir(self._path): + path = (self._path / filename) + if path.is_file() and filename != "index.json": + if (time() - getmtime(path) > self._lifetime): + self._logger.debug(f"Removing {filename}") + to_delete = filename.split('_')[0] + try: + del self._cache[to_delete] + except KeyError: + self._logger.warning(f"{to_delete} has been already removed from the index") + remove(path) + with open(self._path / "index.json", "w") as f: + dump(self._cache, f) + + self._logger.debug("Finished collecting old cache") + + await sleep(self._gc_delay) diff --git a/src/petthecord/petter/exceptions.py b/src/petthecord/petter/exceptions.py new file mode 100644 index 0000000..fcce96e --- /dev/null +++ b/src/petthecord/petter/exceptions.py @@ -0,0 +1,10 @@ +class NotFound(Exception): + pass + + +class HTTPException(Exception): + pass + + +class CacheEnvironmentFail(Exception): + pass diff --git a/src/petthecord/petter/petter.py b/src/petthecord/petter/petter.py new file mode 100644 index 0000000..3e067e4 --- /dev/null +++ b/src/petthecord/petter/petter.py @@ -0,0 +1,45 @@ +from io import BytesIO +from logging import getLogger + +import discord +from discord import Client, User +from petpetgif import petpet + +from .exceptions import HTTPException, NotFound + + +class Petter: + _client: Client + + def __init__(self) -> None: + self._logger = getLogger(__name__) + + def prepare(self, client: Client) -> None: + self._client = client + + async def _get_user(self, uid: int) -> User: + try: + user = await self._client.fetch_user(uid) + return user + except discord.NotFound: + raise NotFound + except discord.HTTPException: + raise HTTPException + + async def _get_avatar(self, user: User) -> bytes: + if not user.avatar: + raise NotFound + return await user.avatar.read() + + async def petpet(self, user: int | User) -> bytes: + if isinstance(user, int): + _user = await self._get_user(user) + else: + _user = user + + self._logger.debug(f"Generating gif for {_user.id}") + + with BytesIO() as f: + petpet.make(BytesIO(await self._get_avatar(_user)), f) + f.seek(0) + return f.read() diff --git a/src/petthecord/runner.py b/src/petthecord/runner.py index 75ed6ca..f3ea325 100644 --- a/src/petthecord/runner.py +++ b/src/petthecord/runner.py @@ -5,7 +5,7 @@ from discord import Intents from discord.ext import commands from .bot import PetTheCordCog -from .cache import CachedPet +from .petter import CachedPetter, Petter from .server import Server @@ -16,10 +16,7 @@ class PetTheCord(commands.AutoShardedBot): host: str = "127.0.0.1", port: int = 8080, origin: str = "https://ptc.pwn3t.ru", - caching: bool = True, - cache_path: str = "/var/cache/petthecord", - cache_lifetime: int = 86400, - cache_gc_delay: int = 14400, + petter: Petter | None = None, shard_count: int = 1, ) -> None: super().__init__( @@ -30,30 +27,22 @@ class PetTheCord(commands.AutoShardedBot): self._host = host self._port = port self._origin = origin - self._caching = caching - self._cache_path = cache_path - self._cache_lifetime = cache_lifetime - self._cache_gc_delay = cache_gc_delay + self._petter = petter or Petter() self._logger = getLogger(__name__) + self._petter.prepare(self) + async def on_ready(self) -> None: await self.add_cog(PetTheCordCog(self._origin)) await self.tree.sync() - petter = CachedPet( - self, - self._caching, - self._cache_path, - self._cache_lifetime, - self._cache_gc_delay - ) - runner = AppRunner(Server(petter)) + runner = AppRunner(Server(self._petter)) await runner.setup() site = TCPSite(runner, self._host, self._port) await site.start() self._logger.info(f"Started serving on {self._host}:{self._port}") - if self._caching: - await petter.gc_loop() + if isinstance(self._petter, CachedPetter): + await self._petter.gc_loop() diff --git a/src/petthecord/server.py b/src/petthecord/server.py index 08d7389..c1f7f8b 100644 --- a/src/petthecord/server.py +++ b/src/petthecord/server.py @@ -3,11 +3,11 @@ from typing import NoReturn from aiohttp.web import Application, StreamResponse, get, HTTPFound, Request, Response -from .cache import CachedPet, HTTPException, NotFound +from .petter import Petter, HTTPException, NotFound class Server(Application): - def __init__(self, petter: CachedPet) -> None: + def __init__(self, petter: Petter) -> None: self._petter = petter super().__init__() |