66 lines
1.7 KiB
Python
66 lines
1.7 KiB
Python
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",
|
|
)
|