# Copyright (c) Streamlit Inc. (2018-2022) Snowflake Inc. (2022-2026) # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. from __future__ import annotations import copy import json from collections.abc import Mapping from dataclasses import dataclass from enum import Enum from typing import ( TYPE_CHECKING, Any, Final, Literal, TypeAlias, overload, ) from streamlit.dataframe_util import DataFormat from streamlit.elements.lib.column_types import ( ButtonColumnResult, ColumnConfig, ColumnType, ) from streamlit.elements.lib.dicttools import remove_none_values from streamlit.elements.lib.policies import check_widget_policies from streamlit.elements.lib.utils import compute_and_register_element_id from streamlit.errors import ( StreamlitAPIException, StreamlitInvalidParameterTypeError, ) from streamlit.runtime.state import register_widget, validate_on_change_mode from streamlit.util import ReadOnlyAttributeDictionary if TYPE_CHECKING: import pyarrow as pa from pandas import DataFrame, Index, Series from streamlit.delta_generator import DeltaGenerator from streamlit.proto.Common_pb2 import StringTriggerValue from streamlit.proto.Dataframe_pb2 import Dataframe as DataframeProto from streamlit.runtime.scriptrunner_utils.script_run_context import ScriptRunContext # The index identifier can be used to apply configuration options IndexIdentifierType = Literal["_index"] INDEX_IDENTIFIER: IndexIdentifierType = "_index" # This is used as prefix for columns that are configured via the numerical position. # The integer value is converted into a string key with this prefix. # This needs to match with the prefix configured in the frontend. NUMERICAL_POSITION_PREFIX = "_pos:" # The column data kind is used to describe the type of the data within the column. class ColumnDataKind(str, Enum): INTEGER = "integer" FLOAT = "float" DATE = "date" TIME = "time" DATETIME = "datetime" BOOLEAN = "boolean" STRING = "string" TIMEDELTA = "timedelta" PERIOD = "period" INTERVAL = "interval" BYTES = "bytes" DECIMAL = "decimal" COMPLEX = "complex" LIST = "list" DICT = "dict" EMPTY = "empty" UNKNOWN = "unknown" # The dataframe schema is a mapping from the name of the column # in the underlying dataframe to the column data kind. # The index column uses `_index` as name. DataframeSchema: TypeAlias = dict[str, ColumnDataKind] # This mapping contains all editable column types mapped to the data kinds # that the column type is compatible for editing. _EDITING_COMPATIBILITY_MAPPING: Final[dict[ColumnType, list[ColumnDataKind]]] = { "text": [ColumnDataKind.STRING, ColumnDataKind.EMPTY], "number": [ ColumnDataKind.INTEGER, ColumnDataKind.FLOAT, ColumnDataKind.DECIMAL, ColumnDataKind.STRING, ColumnDataKind.TIMEDELTA, ColumnDataKind.EMPTY, ], "checkbox": [ ColumnDataKind.BOOLEAN, ColumnDataKind.STRING, ColumnDataKind.INTEGER, ColumnDataKind.EMPTY, ], "selectbox": [ ColumnDataKind.STRING, ColumnDataKind.BOOLEAN, ColumnDataKind.INTEGER, ColumnDataKind.FLOAT, ColumnDataKind.EMPTY, ], "date": [ColumnDataKind.DATE, ColumnDataKind.DATETIME, ColumnDataKind.EMPTY], "time": [ColumnDataKind.TIME, ColumnDataKind.DATETIME, ColumnDataKind.EMPTY], "datetime": [ ColumnDataKind.DATETIME, ColumnDataKind.DATE, ColumnDataKind.TIME, ColumnDataKind.EMPTY, ], "link": [ColumnDataKind.STRING, ColumnDataKind.EMPTY], "list": [ ColumnDataKind.LIST, ColumnDataKind.STRING, ColumnDataKind.EMPTY, ], "multiselect": [ ColumnDataKind.LIST, ColumnDataKind.STRING, ColumnDataKind.EMPTY, ], "markdown": [ColumnDataKind.STRING, ColumnDataKind.EMPTY], } def is_type_compatible(column_type: ColumnType, data_kind: ColumnDataKind) -> bool: """Check if the column type is compatible with the underlying data kind. This check only applies to editable column types (e.g. number or text). Non-editable column types (e.g. bar_chart or image) can be configured for all data kinds (this might change in the future). Parameters ---------- column_type : ColumnType The column type to check. data_kind : ColumnDataKind The data kind to check. Returns ------- bool True if the column type is compatible with the data kind, False otherwise. """ if column_type not in _EDITING_COMPATIBILITY_MAPPING: return True return data_kind in _EDITING_COMPATIBILITY_MAPPING[column_type] def _determine_data_kind_via_arrow(field: pa.Field) -> ColumnDataKind: """Determine the data kind via the arrow type information. The column data kind refers to the shared data type of the values in the column (e.g. int, float, str, bool). Parameters ---------- field : pa.Field The arrow field from the arrow table schema. Returns ------- ColumnDataKind The data kind of the field. """ import pyarrow as pa field_type = field.type if pa.types.is_integer(field_type): return ColumnDataKind.INTEGER if pa.types.is_floating(field_type): return ColumnDataKind.FLOAT if pa.types.is_boolean(field_type): return ColumnDataKind.BOOLEAN if pa.types.is_string(field_type) or pa.types.is_large_string(field_type): return ColumnDataKind.STRING if pa.types.is_date(field_type): return ColumnDataKind.DATE if pa.types.is_time(field_type): return ColumnDataKind.TIME if pa.types.is_timestamp(field_type): return ColumnDataKind.DATETIME if pa.types.is_duration(field_type): return ColumnDataKind.TIMEDELTA if pa.types.is_list(field_type): return ColumnDataKind.LIST if pa.types.is_decimal(field_type): return ColumnDataKind.DECIMAL if pa.types.is_null(field_type): return ColumnDataKind.EMPTY # Interval does not seem to work correctly: # if pa.types.is_interval(field_type): # return ColumnDataKind.INTERVAL # noqa: ERA001 if pa.types.is_binary(field_type): return ColumnDataKind.BYTES if pa.types.is_struct(field_type): return ColumnDataKind.DICT return ColumnDataKind.UNKNOWN def _determine_data_kind_via_pandas_dtype( column: Series[Any] | Index[Any], ) -> ColumnDataKind: """Determine the data kind by using the pandas dtype. The column data kind refers to the shared data type of the values in the column (e.g. int, float, str, bool). Parameters ---------- column : pd.Series, pd.Index The column for which the data kind should be determined. Returns ------- ColumnDataKind The data kind of the column. """ import pandas as pd column_dtype = column.dtype if pd.api.types.is_bool_dtype(column_dtype): return ColumnDataKind.BOOLEAN if pd.api.types.is_integer_dtype(column_dtype): return ColumnDataKind.INTEGER if pd.api.types.is_float_dtype(column_dtype): return ColumnDataKind.FLOAT if pd.api.types.is_datetime64_any_dtype(column_dtype): return ColumnDataKind.DATETIME if pd.api.types.is_timedelta64_dtype(column_dtype): return ColumnDataKind.TIMEDELTA if isinstance(column_dtype, pd.PeriodDtype): return ColumnDataKind.PERIOD if isinstance(column_dtype, pd.IntervalDtype): return ColumnDataKind.INTERVAL if pd.api.types.is_complex_dtype(column_dtype): return ColumnDataKind.COMPLEX if pd.api.types.is_object_dtype( column_dtype ) is False and pd.api.types.is_string_dtype(column_dtype): # This handles pandas 3.0+ StringDtype (and PyArrow-backed string types). # We exclude object dtype here because object columns with string values # are handled via _determine_data_kind_via_inferred_type in the caller. return ColumnDataKind.STRING return ColumnDataKind.UNKNOWN def _determine_data_kind_via_inferred_type( column: Series[Any] | Index[Any], ) -> ColumnDataKind: """Determine the data kind by inferring it from the underlying data. The column data kind refers to the shared data type of the values in the column (e.g. int, float, str, bool). Parameters ---------- column : pd.Series, pd.Index The column to determine the data kind for. Returns ------- ColumnDataKind The data kind of the column. """ from pandas.api.types import infer_dtype inferred_type = infer_dtype(column) if inferred_type == "string": return ColumnDataKind.STRING if inferred_type == "bytes": return ColumnDataKind.BYTES if inferred_type in {"floating", "mixed-integer-float"}: return ColumnDataKind.FLOAT if inferred_type == "integer": return ColumnDataKind.INTEGER if inferred_type == "decimal": return ColumnDataKind.DECIMAL if inferred_type == "complex": return ColumnDataKind.COMPLEX if inferred_type == "boolean": return ColumnDataKind.BOOLEAN if inferred_type in {"datetime64", "datetime"}: return ColumnDataKind.DATETIME if inferred_type == "date": return ColumnDataKind.DATE if inferred_type in {"timedelta64", "timedelta"}: return ColumnDataKind.TIMEDELTA if inferred_type == "time": return ColumnDataKind.TIME if inferred_type == "period": return ColumnDataKind.PERIOD if inferred_type == "interval": return ColumnDataKind.INTERVAL if inferred_type == "empty": return ColumnDataKind.EMPTY # Unused types: mixed, unknown-array, categorical, mixed-integer return ColumnDataKind.UNKNOWN def _determine_data_kind( column: Series[Any] | Index[Any], field: pa.Field | None = None ) -> ColumnDataKind: """Determine the data kind of a column. The column data kind refers to the shared data type of the values in the column (e.g. int, float, str, bool). Parameters ---------- column : pd.Series, pd.Index The column to determine the data kind for. field : pa.Field, optional The arrow field from the arrow table schema. Returns ------- ColumnDataKind The data kind of the column. """ import pandas as pd if isinstance(column.dtype, pd.CategoricalDtype): # Categorical columns can have different underlying data kinds # depending on the categories. return _determine_data_kind_via_inferred_type(column.dtype.categories) if field is not None: data_kind = _determine_data_kind_via_arrow(field) if data_kind != ColumnDataKind.UNKNOWN: return data_kind if column.dtype.name == "object": # If dtype is object, we need to infer the type from the column return _determine_data_kind_via_inferred_type(column) return _determine_data_kind_via_pandas_dtype(column) def determine_dataframe_schema( data_df: DataFrame, arrow_schema: pa.Schema ) -> DataframeSchema: """Determine the schema of a dataframe. Parameters ---------- data_df : pd.DataFrame The dataframe to determine the schema of. arrow_schema : pa.Schema The Arrow schema of the dataframe. Returns ------- DataframeSchema A mapping that contains the detected data type for the index and columns. The key is the column name in the underlying dataframe or ``_index`` for index columns. """ dataframe_schema: DataframeSchema = {} # Add type of index: # TODO(lukasmasuch): We need to apply changes here to support multiindex. dataframe_schema[INDEX_IDENTIFIER] = _determine_data_kind(data_df.index) # Add types for all columns: for i, column in enumerate(data_df.items()): column_name = str(column[0]) column_data = column[1] dataframe_schema[column_name] = _determine_data_kind( column_data, arrow_schema.field(i) ) return dataframe_schema # A mapping of column names/IDs to column configs. ColumnConfigMapping: TypeAlias = dict[IndexIdentifierType | str | int, ColumnConfig] ColumnConfigMappingInput: TypeAlias = Mapping[ # TODO(lukasmasuch): This should also use int here to # correctly type the support for positional index. However, # allowing int here leads mypy to complain about simple dict[str, ...] # as input -> which seems like a mypy bug. IndexIdentifierType | str, ColumnConfig | ButtonColumnResult | str | None, ] ButtonColumnMapping: TypeAlias = dict[str, ButtonColumnResult] class ButtonColumnClickState(ReadOnlyAttributeDictionary): """The schema for button click state in ButtonColumn. To use this type in an annotation, import it from ``streamlit.typing``. Read-only dict-like click payload with attribute and key access (``click.row`` / ``click["row"]``). Both fields are always present when a button click occurs. """ row: int label: str @overload def __getitem__(self, key: Literal["row"]) -> int: ... @overload def __getitem__(self, key: Literal["label"]) -> str: ... @overload def __getitem__(self, key: Any) -> Any: ... def __getitem__(self, key: Any) -> Any: return super().__getitem__(key) @dataclass class ButtonClickSerde: """Serializer/deserializer for ButtonColumn click values. Uses string trigger value pattern (value resets after each run). The frontend sends the click state as a JSON string. """ def serialize(self, v: ButtonColumnClickState | None) -> StringTriggerValue: from streamlit.proto.Common_pb2 import StringTriggerValue if v is None: return StringTriggerValue() return StringTriggerValue(data=json.dumps(v)) def deserialize(self, ui_value: str | None) -> ButtonColumnClickState | None: if not ui_value: return None parsed = json.loads(ui_value) # Validate shape: must be a dict with "row" (int >= 0) and "label" (str). # bool is a subclass of int, so booleans are explicitly rejected for "row". if ( not isinstance(parsed, dict) or not isinstance(parsed.get("row"), int) or isinstance(parsed.get("row"), bool) or not isinstance(parsed.get("label"), str) ): raise StreamlitAPIException( "Invalid button click state: expected {row: int, label: str}.", error_id="button-column-invalid-click-state", ) # Validate row is non-negative (bounds check - row < num_rows is # checked downstream when accessing the data) if parsed["row"] < 0: raise StreamlitAPIException( f"Invalid button click row index: {parsed['row']}. Row must be >= 0.", error_id="button-column-invalid-click-row", ) return ButtonColumnClickState(parsed) def extract_button_column_configs( column_config: ColumnConfigMappingInput | None, ) -> tuple[dict[Any, ColumnConfig | str | None] | None, ButtonColumnMapping]: """Extract interactive button-column configs from a column config mapping.""" button_columns: ButtonColumnMapping = {} processed_column_config: dict[Any, ColumnConfig | str | None] | None = None if column_config is not None: processed_column_config = {} for col_name, config in column_config.items(): if isinstance(config, ButtonColumnResult): # Transform key the same way column config does for consistency. column_widget_key = ( f"{NUMERICAL_POSITION_PREFIX}{col_name}" if isinstance(col_name, int) else str(col_name) ) button_columns[column_widget_key] = config processed_column_config[col_name] = config.config else: processed_column_config[col_name] = config return processed_column_config, button_columns def register_button_column_widgets( *, dg: DeltaGenerator, proto: DataframeProto, button_columns: ButtonColumnMapping, ctx: ScriptRunContext | None, ) -> None: """Register widgets for interactive button columns and attach them to a dataframe proto.""" button_serde = ButtonClickSerde() for col_name, button_col in button_columns.items(): on_click = validate_on_change_mode( button_col.on_click, supported_modes=(), param_name="on_click", ) check_widget_policies( dg, button_col.key, on_change=on_click, default_value=None, writes_allowed=False, ) widget_id = compute_and_register_element_id( "dataframe_button", user_key=button_col.key, key_as_main_identity=True, dg=dg, ) register_widget( widget_id, on_change_handler=on_click, args=button_col.args, kwargs=button_col.kwargs, deserializer=button_serde.deserialize, serializer=button_serde.serialize, ctx=ctx, value_type="string_trigger_value", ) proto.button_click_widgets[col_name] = widget_id def process_config_mapping( column_config: ColumnConfigMappingInput | None = None, ) -> ColumnConfigMapping: """Transforms a user-provided column config mapping into a valid column config mapping that can be used by the frontend. Parameters ---------- column_config: dict or None The user-provided column config mapping. Returns ------- dict The transformed column config mapping. """ if column_config is None: return {} transformed_column_config: ColumnConfigMapping = {} for column, config in column_config.items(): if config is None: transformed_column_config[column] = ColumnConfig(hidden=True) elif isinstance(config, str): transformed_column_config[column] = ColumnConfig(label=config) elif isinstance(config, ButtonColumnResult): # ButtonColumnResult is typically preprocessed before reaching this function # in both st.dataframe and st.data_editor. If we encounter it here, extract # the config dict. Button columns are always read-only. transformed_column_config[column] = copy.deepcopy(config.config) elif isinstance(config, dict): # Ensure that the column config objects are cloned # since we will apply in-place changes to it. transformed_column_config[column] = copy.deepcopy(config) else: raise StreamlitInvalidParameterTypeError( "column_config", type(config).__name__, ["None", "str", "dict"], detail=f"Invalid configuration for column `{column}`.", ) return transformed_column_config def update_column_config( column_config_mapping: ColumnConfigMapping, column: str | int, column_config: ColumnConfig, ) -> None: """Updates the column config value for a single column within the mapping. Parameters ---------- column_config_mapping : ColumnConfigMapping The column config mapping to update. column : str | int The column to update the config value for. This can be the column name or the numerical position of the column. column_config : ColumnConfig The column config to update. """ if column not in column_config_mapping: column_config_mapping[column] = {} column_config_mapping[column].update(column_config) def apply_data_specific_configs( columns_config: ColumnConfigMapping, data_format: DataFormat, ) -> None: """Apply data specific configurations to the provided dataframe. This will apply inplace changes to the dataframe and the column configurations depending on the data format. Parameters ---------- columns_config : ColumnConfigMapping A mapping of column names/ids to column configurations. data_format : DataFormat The format of the data. """ # Pandas adds a range index as default to all datastructures # but for most of the non-pandas data objects it is unnecessary # to show this index to the user. Therefore, we will hide it as default. if data_format in { DataFormat.SET_OF_VALUES, DataFormat.TUPLE_OF_VALUES, DataFormat.LIST_OF_VALUES, DataFormat.NUMPY_LIST, DataFormat.NUMPY_MATRIX, DataFormat.LIST_OF_RECORDS, DataFormat.LIST_OF_ROWS, DataFormat.COLUMN_VALUE_MAPPING, # Dataframe-like objects that don't have an index: DataFormat.PANDAS_ARRAY, DataFormat.PANDAS_INDEX, DataFormat.POLARS_DATAFRAME, DataFormat.POLARS_SERIES, DataFormat.POLARS_LAZYFRAME, DataFormat.PYARROW_ARRAY, }: update_column_config(columns_config, INDEX_IDENTIFIER, {"hidden": True}) def _convert_column_config_to_json(column_config_mapping: ColumnConfigMapping) -> str: try: # Ignore all None values and prefix columns specified by numerical index: return json.dumps( { (f"{NUMERICAL_POSITION_PREFIX}{k!s}" if isinstance(k, int) else k): v for (k, v) in remove_none_values(column_config_mapping).items() }, allow_nan=False, ) except ValueError as ex: raise StreamlitAPIException( f"The provided column config cannot be serialized into JSON: {ex}", error_id="column-config-json-serialize-failed", ) from ex def marshall_column_config( proto: DataframeProto, column_config_mapping: ColumnConfigMapping ) -> None: """Marshall the column config into the Dataframe proto. Parameters ---------- proto : DataframeProto The proto to marshall into. column_config_mapping : ColumnConfigMapping The column config to marshall. """ proto.columns = _convert_column_config_to_json(column_config_mapping)