from __future__ import annotations
from typing import TYPE_CHECKING

from typing import Type, Union, Literal, Optional, Dict, List, Tuple, Set, Annotated, Any
from collections.abc import AsyncIterator
from .common import CommonDatasetsProvider, ComponentActionContext

if TYPE_CHECKING:
    from datasets import Dataset

class LocalDatasetsProvider(CommonDatasetsProvider):
    async def load(self, context: ComponentActionContext) -> Dataset:
        from datasets import load_dataset

        loader         = await context.render_variable(self.config.loader)
        data_files     = await context.render_variable(self.config.data_files)
        data_dir       = await context.render_variable(self.config.data_dir)
        split          = await context.render_variable(self.config.split)
        streaming      = await context.render_variable(self.config.streaming)
        keep_in_memory = await context.render_variable(self.config.keep_in_memory)
        cache_dir      = await context.render_variable(self.config.cache_dir)
        save_infos     = await context.render_variable(self.config.save_infos)

        return load_dataset(
            path=loader,
            data_files=data_files,
            data_dir=data_dir,
            split=split,
            streaming=streaming,
            keep_in_memory=keep_in_memory,
            cache_dir=cache_dir,
            save_infos=save_infos
        )
