from __future__ import annotations

import typing as t

import pyarrow as pa

from sarus_synthetic_data.configs.global_config import (
    TableConfig,
)
from sarus_synthetic_data.correlations_generator import typing as syn_typing
import sarus_synthetic_data.constants as syn_constant


class PreprocessingConfig(t.Protocol):
    saving_directory: str
    tables: t.Mapping[t.Tuple[str, ...], TableConfig]
    privacy_unit_col: str


class OptionalOutput(t.TypedDict):
    syn_constant.IS_NOT_NULL: pa.Array  # type:ignore #noqa:B032
    syn_constant.OPTIONAL_VALUE: syn_typing.PyTree  # type:ignore #noqa:B032