فهرست منبع

Add Lazy injector wrapper

Kirk Trombley 4 سال پیش
والد
کامیت
0c1603d702
3فایلهای تغییر یافته به همراه17 افزوده شده و 3 حذف شده
  1. 1 1
      lib/rollbot/command/__init__.py
  2. 14 0
      lib/rollbot/command/injection.py
  3. 2 2
      repl_driver.py

+ 1 - 1
lib/rollbot/command/__init__.py

@@ -1,3 +1,3 @@
 from .decorators import as_command, on_startup, on_shutdown, get_command_config
 from .failure import RollbotFailure
-from .injection import Args, ArgList, ArgListSplitOn, ArgParse, Database
+from .injection import Args, ArgList, ArgListSplitOn, ArgParse, Database, Lazy

+ 14 - 0
lib/rollbot/command/injection.py

@@ -1,5 +1,6 @@
 from typing import Generic, TypeVar, Optional
 from argparse import ArgumentParser, Namespace
+from collections.abc import Callable, Coroutine
 import shlex
 
 from aiosqlite.core import Connection
@@ -52,6 +53,19 @@ class DatabaseInjector(Injector[Connection]):
         return context.database()
 
 
+class Lazy(Injector[Callable[[], Coroutine[None, None, Dep]]]):
+    def __init__(self, deferred: Injector[Dep]):
+        self.deferred = deferred
+        self._calculated = None
+
+    async def inject(self, message: Message, context: Context) -> Callable[[], Coroutine[None, None, Dep]]:
+        async def call():
+            if self._calculated is None:
+                self._calculated = await self.deferred.inject(message, context)
+            return self._calculated
+        return call
+
+
 Args = ArgsInjector()
 ArgList = ArgListSplitOn()
 Database = DatabaseInjector()

+ 2 - 2
repl_driver.py

@@ -46,9 +46,9 @@ async def count_command(message, context):
 
 
 @rollbot.as_command
-async def count2(args: rollbot.ArgList, db: rollbot.Database):
+async def count2(args: rollbot.ArgList, connect: rollbot.Lazy(rollbot.Database)):
     name = args[0] if len(args) > 0 else "main"
-    async with db:
+    async with await connect() as db:
         await db.execute(
             "INSERT INTO counter VALUES (?, 1) \
             ON CONFLICT (name) DO UPDATE SET count=count + 2",