Streamline initialization logic

lib/rollbot/decorators/__init__.py

@@ -2,12 +2,14 @@ from ..types import CommandConfiguration
 
 from .as_command import as_command, decorated_commands
 from .on_event import on_startup, on_shutdown, decorated_startup, decorated_shutdown
+from .initialize_data import initialize_data
 from .error_handling import with_failure_handling
 
 __all__ = [
     "as_command",
     "on_startup",
     "on_shutdown",
+    "initialize_data",
     "with_failure_handling",
     "get_command_config",
 ]

lib/rollbot/decorators/initialize_data.py

@@ -0,0 +1,18 @@
+from typing import Type, TypeVar
+
+from ..types import Context
+from ..injection import Data
+
+from .on_event import on_startup
+
+
+T = TypeVar("T")
+
+
+def initialize_data(cls: Type[T]) -> Type[T]:
+    @on_startup
+    async def setup_my_counter(context: Context):
+        async with context.database() as db:
+            await Data.initialize(cls, db)
+
+    return cls

lib/rollbot/injection/data.py

@@ -37,7 +37,7 @@ class DataStore(Generic[DataType]):
             ("_" + c.lower()) if "A" <= c <= "Z" else c for c in datatype.__name__ if c.isalnum()
         ).strip("_")
 
-    async def setup(self):
+    async def _setup(self):
         await self.connection.execute(
             f'CREATE TABLE IF NOT EXISTS {self.table_name} ( \
                 key TEXT NOT NULL PRIMARY KEY, \
@@ -88,7 +88,6 @@ class DataFor(Injector[Optional[DataType]]):
         key = await self.key.inject(message, context)
         async with context.database() as db:
             store = DataStore(self.datatype, db)
-            await store.setup()
             return await store.load_or(key, **self.kwargs)
 
 
@@ -97,9 +96,13 @@ class Data(InjectorWithCleanup[DataStore[DataType]]):
         self.datatype = datatype
         self.For = lambda key, **kw: DataFor(datatype, key, kw)
 
+    @staticmethod
+    async def initialize(datatype: Type[DataType], connection: Connection):
+        store = DataStore(datatype, connection)
+        await store._setup()
+
     async def inject(self, message: Message, context: Context) -> DataStore[DataType]:
         store = DataStore(self.datatype, await context.database())
-        await store.setup()
         return store
 
     async def cleanup(self, store: DataStore[DataType]):

repl_driver.py

@@ -6,6 +6,7 @@ import rollbot
 from rollbot.injection import ArgList, Lazy, Database, Data, Args
 
 
+@rollbot.initialize_data
 @dataclass
 class MyCounter:
     one: int = 0