From 09fe2c1160cf87d31ebe381d1622d5889c8a333d Mon Sep 17 00:00:00 2001
From: Nakidai <nakidai@disroot.org>
Date: Tue, 5 Nov 2024 16:22:44 +0300
Subject: Add shard support, fix #2

Shards are way to fix the problem with 429
---
 src/petthecord/main.py   | 10 +++++++++-
 src/petthecord/runner.py |  6 ++++--
 2 files changed, 13 insertions(+), 3 deletions(-)

diff --git a/src/petthecord/main.py b/src/petthecord/main.py
index c164685..9177926 100644
--- a/src/petthecord/main.py
+++ b/src/petthecord/main.py
@@ -55,6 +55,13 @@ def main() -> None:
         metavar="TIME",
         help="Delay between cache's garbage collector runs in seconds"
     )
+    parser.add_argument(
+        "-c", "--shards",
+        default=1,
+        type=int,
+        metavar="COUNT",
+        help="Amount of shards to create"
+    )
     args = parser.parse_args()
 
     bot = PetTheCord(
@@ -64,7 +71,8 @@ def main() -> None:
         not args.no_cache,
         args.cache_dir,
         args.cache_lifetime,
-        args.cache_gc_delay
+        args.cache_gc_delay,
+        args.shards,
     )
     if (token := getenv("PETTHECORD_TOKEN")) is not None:
         bot.run(token, root_logger=True)
diff --git a/src/petthecord/runner.py b/src/petthecord/runner.py
index f40577c..75ed6ca 100644
--- a/src/petthecord/runner.py
+++ b/src/petthecord/runner.py
@@ -9,7 +9,7 @@ from .cache import CachedPet
 from .server import Server
 
 
-class PetTheCord(commands.Bot):
+class PetTheCord(commands.AutoShardedBot):
     def __init__(
         self,
 
@@ -20,10 +20,12 @@ class PetTheCord(commands.Bot):
         cache_path: str = "/var/cache/petthecord",
         cache_lifetime: int = 86400,
         cache_gc_delay: int = 14400,
+        shard_count: int = 1,
     ) -> None:
         super().__init__(
             command_prefix="!",
-            intents=Intents.default()
+            intents=Intents.default(),
+            shard_count=shard_count,
         )
         self._host = host
         self._port = port
-- 
cgit 1.4.1