about summary refs log tree commit diff
diff options
context:
space:
mode:
authorNakidai <nakidai@disroot.org>2024-11-07 03:02:00 +0300
committerNakidai <nakidai@disroot.org>2024-11-07 03:02:00 +0300
commit4f99e6ada3ddde9c8b1a79129db7653e5862c985 (patch)
tree0c8e5445c23cf23c292b1a4bac8e063d257f0da2
parentdda6671fdbbbd6b3a1f2529dbc60726b67cd4fb2 (diff)
downloadpetthecord-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__.py3
-rw-r--r--src/petthecord/cache.py121
-rw-r--r--src/petthecord/main.py15
-rw-r--r--src/petthecord/petter/__init__.py6
-rw-r--r--src/petthecord/petter/cachedpetter.py97
-rw-r--r--src/petthecord/petter/exceptions.py10
-rw-r--r--src/petthecord/petter/petter.py45
-rw-r--r--src/petthecord/runner.py27
-rw-r--r--src/petthecord/server.py4
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__()