Source code for rics.translation.dio._resolve

from typing import Dict, Sequence, Type, TypeVar

import pandas as pd

from rics.translation.dio import DataStructureIO
from rics.translation.dio._dict import DictIO
from rics.translation.dio._pandas import PandasIO
from rics.translation.dio._sequence import SequenceIO
from rics.translation.dio._single_value import SingleValueIO
from rics.translation.dio.exceptions import UntranslatableTypeError

DefaultTranslatable = TypeVar("DefaultTranslatable", pd.DataFrame, pd.Series, Dict, Sequence, str, int)


def resolve_io(arg: DefaultTranslatable) -> Type[DataStructureIO]:
    """Get an IO instance for `arg`.

    Args:
        arg: An argument to get IO for.

    Returns:
        A data structure IO instance for `arg`.

    Raises:
        UntranslatableTypeError: If not IO could be found.
    """
    for tio_class in (PandasIO, DictIO, SequenceIO, SingleValueIO):
        if tio_class.handles_type(arg):
            return tio_class

    raise UntranslatableTypeError(type(arg))