from typing import Type, TypeVar
import dataclasses

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]:
    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:
            await Data.initialize(cls, db)

    return cls