Bladeren bron

Add auto-dataclass-ification to initialize_data

Kirk Trombley 4 jaren geleden
bovenliggende
commit
4b4e53a3fb
2 gewijzigde bestanden met toevoegingen van 4 en 2 verwijderingen
  1. 4 0
      lib/rollbot/decorators/initialize_data.py
  2. 0 2
      repl_driver.py

+ 4 - 0
lib/rollbot/decorators/initialize_data.py

@@ -1,4 +1,5 @@
 from typing import Type, TypeVar
+import dataclasses
 
 from ..types import Context
 from ..injection import Data
@@ -10,6 +11,9 @@ T = TypeVar("T")
 
 
 def initialize_data(cls: Type[T]) -> Type[T]:
+    if not dataclasses.is_dataclass(cls):
+        cls = dataclasses.dataclass(cls)
+
     @on_startup
     async def setup_my_counter(context: Context):
         async with context.database() as db:

+ 0 - 2
repl_driver.py

@@ -1,5 +1,4 @@
 from datetime import datetime
-from dataclasses import dataclass
 import asyncio
 
 import rollbot
@@ -7,7 +6,6 @@ from rollbot.injection import ArgList, Lazy, Database, Data, Args
 
 
 @rollbot.initialize_data
-@dataclass
 class MyCounter:
     one: int = 0
     two: int = 0