|
from __future__ import annotations |
|
|
|
from typing import TYPE_CHECKING, Callable, overload |
|
|
|
from altair.utils.core import sanitize_pandas_dataframe |
|
from altair.utils.data import DataTransformerRegistry as _DataTransformerRegistry |
|
from altair.utils.data import ( |
|
MaxRowsError, |
|
check_data_type, |
|
limit_rows, |
|
sample, |
|
to_csv, |
|
to_json, |
|
to_values, |
|
) |
|
|
|
if TYPE_CHECKING: |
|
from altair.utils.data import DataType, ToValuesReturnType |
|
from altair.utils.plugin_registry import PluginEnabler |
|
|
|
|
|
@overload |
|
def default_data_transformer( |
|
data: None = ..., max_rows: int = ... |
|
) -> Callable[[DataType], ToValuesReturnType]: ... |
|
@overload |
|
def default_data_transformer( |
|
data: DataType, max_rows: int = ... |
|
) -> ToValuesReturnType: ... |
|
def default_data_transformer( |
|
data: DataType | None = None, max_rows: int = 5000 |
|
) -> Callable[[DataType], ToValuesReturnType] | ToValuesReturnType: |
|
if data is None: |
|
|
|
def pipe(data: DataType, /) -> ToValuesReturnType: |
|
data = limit_rows(data, max_rows=max_rows) |
|
return to_values(data) |
|
|
|
return pipe |
|
|
|
else: |
|
return to_values(limit_rows(data, max_rows=max_rows)) |
|
|
|
|
|
class DataTransformerRegistry(_DataTransformerRegistry): |
|
def disable_max_rows(self) -> PluginEnabler: |
|
"""Disable the MaxRowsError.""" |
|
options = self.options |
|
if self.active in {"default", "vegafusion"}: |
|
options = options.copy() |
|
options["max_rows"] = None |
|
return self.enable(**options) |
|
|
|
|
|
__all__ = ( |
|
"DataTransformerRegistry", |
|
"MaxRowsError", |
|
"check_data_type", |
|
"default_data_transformer", |
|
"limit_rows", |
|
"sample", |
|
"sanitize_pandas_dataframe", |
|
"to_csv", |
|
"to_json", |
|
"to_values", |
|
) |
|
|
