streamlit-nightly 1.33.1.dev20240501__py2.py3-none-any.whl → 1.34.1.dev20240503__py2.py3-none-any.whl

streamlit/elements/{arrow_altair.py → vega_charts.py}

@@ -12,30 +12,22 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""A Python wrapper around Altair.
-Altair is a Python visualization library based on Vega-Lite,
-a nice JSON schema for expressing graphs and charts.
-"""
+"""Collection of chart commands that are rendered via our vega-lite chart component."""
+
 from __future__ import annotations
 
+import json
 from contextlib import nullcontext
-from datetime import date
-from enum import Enum
-from typing import TYPE_CHECKING, Any, Collection, Literal, Sequence, cast
+from typing import TYPE_CHECKING, Any, Final, Literal, Sequence, cast
 
-import streamlit.elements.arrow_vega_lite as arrow_vega_lite
+import streamlit.elements.lib.dicttools as dicttools
 from streamlit import type_util
-from streamlit.color_util import (
-    Color,
-    is_color_like,
-    is_color_tuple_like,
-    is_hex_color_like,
-    to_css_color,
+from streamlit.elements.lib.built_in_chart_utils import (
+    AddRowsMetadata,
+    ChartType,
+    generate_chart,
 )
-from streamlit.elements.altair_utils import AddRowsMetadata
-from streamlit.elements.arrow import Data
-from streamlit.elements.utils import last_index_for_melted_dataframes
-from streamlit.errors import Error, StreamlitAPIException
+from streamlit.errors import StreamlitAPIException
 from streamlit.proto.ArrowVegaLiteChart_pb2 import (
     ArrowVegaLiteChart as ArrowVegaLiteChartProto,
 )
@@ -43,54 +35,171 @@ from streamlit.runtime.metrics_util import gather_metrics
 
 if TYPE_CHECKING:
     import altair as alt
-    import pandas as pd
 
+    from streamlit.color_util import Color
     from streamlit.delta_generator import DeltaGenerator
+    from streamlit.elements.arrow import Data
+
+# See https://vega.github.io/vega-lite/docs/encoding.html
+_CHANNELS: Final = {
+    "x",
+    "y",
+    "x2",
+    "y2",
+    "xError",
+    "xError2",
+    "yError",
+    "yError2",
+    "longitude",
+    "latitude",
+    "color",
+    "opacity",
+    "fillOpacity",
+    "strokeOpacity",
+    "strokeWidth",
+    "size",
+    "shape",
+    "text",
+    "tooltip",
+    "href",
+    "key",
+    "order",
+    "detail",
+    "facet",
+    "row",
+    "column",
+}
+
+
+def _prepare_vega_lite_spec(
+    spec: dict[str, Any] | None = None,
+    use_container_width: bool = False,
+    **kwargs,
+) -> dict[str, Any]:
+    # Support passing no spec arg, but filling it with kwargs.
+    # Example:
+    #   marshall(proto, baz='boz')
+    if spec is None:
+        spec = dict()
+
+    if len(kwargs):
+        # Support passing in kwargs. Example:
+        #   marshall(proto, {foo: 'bar'}, baz='boz')
+        # Merge spec with unflattened kwargs, where kwargs take precedence.
+        # This only works for string keys, but kwarg keys are strings anyways.
+        spec = dict(spec, **dicttools.unflatten(kwargs, _CHANNELS))
+    else:
+        # Clone the spec dict, since we may be mutating it.
+        spec = dict(spec)
+
+    if len(spec) == 0:
+        raise StreamlitAPIException("Vega-Lite charts require a non-empty spec dict.")
+
+    if "autosize" not in spec:
+        # type fit does not work for many chart types. This change focuses
+        # on vconcat with use_container_width=True as there are unintended
+        # consequences of changing the default autosize for all charts.
+        # fit-x fits the width and height can be adjusted.
+        if "vconcat" in spec and use_container_width:
+            spec["autosize"] = {"type": "fit-x", "contains": "padding"}
+        else:
+            spec["autosize"] = {"type": "fit", "contains": "padding"}
 
+    return spec
 
-class ChartType(Enum):
-    AREA = {"mark_type": "area"}
-    BAR = {"mark_type": "bar"}
-    LINE = {"mark_type": "line"}
-    SCATTER = {"mark_type": "circle"}
 
+def _serialize_data(data: Any) -> bytes:
+    """Serialize the any type of data structure to Arrow IPC format (bytes)."""
+    import pyarrow as pa
+
+    if isinstance(data, pa.Table):
+        return type_util.pyarrow_table_to_bytes(data)
+
+    df = type_util.convert_anything_to_df(data)
+    return type_util.data_frame_to_bytes(df)
+
+
+def _marshall_chart_data(
+    proto: ArrowVegaLiteChartProto,
+    spec: dict[str, Any],
+    data: Data = None,
+) -> None:
+    """Adds the data to the proto and removes it from the spec dict.
+    These operations will happen in-place."""
+
+    # Pull data out of spec dict when it's in a 'datasets' key:
+    #   datasets: {foo: df1, bar: df2}, ...}
+    if "datasets" in spec:
+        for dataset_name, dataset_data in spec["datasets"].items():
+            dataset = proto.datasets.add()
+            dataset.name = str(dataset_name)
+            dataset.has_name = True
+            dataset.data.data = _serialize_data(dataset_data)
+        del spec["datasets"]
+
+    # Pull data out of spec dict when it's in a top-level 'data' key:
+    #   {data: df}
+    #   {data: {values: df, ...}}
+    #   {data: {url: 'url'}}
+    #   {data: {name: 'foo'}}
+    if "data" in spec:
+        data_spec = spec["data"]
+
+        if isinstance(data_spec, dict):
+            if "values" in data_spec:
+                data = data_spec["values"]
+                del spec["data"]
+        else:
+            data = data_spec
+            del spec["data"]
+
+    if data is not None:
+        proto.data.data = _serialize_data(data)
+
+
+def _convert_altair_to_vega_lite_spec(altair_chart: alt.Chart) -> dict[str, Any]:
+    """Convert an Altair chart object to a Vega-Lite chart spec."""
+    import altair as alt
+
+    # Normally altair_chart.to_dict() would transform the dataframe used by the
+    # chart into an array of dictionaries. To avoid that, we install a
+    # transformer that replaces datasets with a reference by the object id of
+    # the dataframe. We then fill in the dataset manually later on.
+
+    datasets = {}
+
+    def id_transform(data) -> dict[str, str]:
+        """Altair data transformer that returns a fake named dataset with the
+        object id.
+        """
+        name = str(id(data))
+        datasets[name] = data
+        return {"name": name}
+
+    alt.data_transformers.register("id", id_transform)  # type: ignore[attr-defined,unused-ignore]
+
+    # The default altair theme has some width/height defaults defined
+    # which are not useful for Streamlit. Therefore, we change the theme to
+    # "none" to avoid those defaults.
+    with alt.themes.enable("none") if alt.themes.active == "default" else nullcontext():  # type: ignore[attr-defined,unused-ignore]
+        with alt.data_transformers.enable("id"):  # type: ignore[attr-defined,unused-ignore]
+            chart_dict = altair_chart.to_dict()
+
+    # Put datasets back into the chart dict but note how they weren't
+    # transformed.
+    chart_dict["datasets"] = datasets
+    return chart_dict
+
+
+class VegaChartsMixin:
+    """Mix-in class for all vega-related chart commands.
+
+    Altair is a python wrapper on top of the vega-lite spec. And our
+    built-in chart commands are just another layer on-top of Altair.
+    All of these chart commands will be eventually converted to a vega-lite
+    spec and rendered using the same vega-lite chart component.
+    """
 
-# Color and size legends need different title paddings in order for them
-# to be vertically aligned.
-#
-# NOTE: I don't think it's possible to *perfectly* align the size and
-# color legends in all instances, since the "size" circles vary in size based
-# on the data, and their container is top-aligned with the color container. But
-# through trial-and-error I found this value to be a good enough middle ground.
-# See e2e/scripts/st_arrow_scatter_chart.py for some alignment tests.
-#
-# NOTE #2: In theory, we could move COLOR_LEGEND_SETTINGS into
-# ArrowVegaLiteChart/CustomTheme.tsx, but this would impact existing behavior.
-# (See https://github.com/streamlit/streamlit/pull/7164#discussion_r1307707345)
-COLOR_LEGEND_SETTINGS = dict(titlePadding=5, offset=5, orient="bottom")
-SIZE_LEGEND_SETTINGS = dict(titlePadding=0.5, offset=5, orient="bottom")
-
-# User-readable names to give the index and melted columns.
-SEPARATED_INDEX_COLUMN_TITLE = "index"
-MELTED_Y_COLUMN_TITLE = "value"
-MELTED_COLOR_COLUMN_TITLE = "color"
-
-# Crazy internal (non-user-visible) names for the index and melted columns, in order to
-# avoid collision with existing column names. The suffix below was generated with an
-# online random number generator. Rationale: because it makes it even less likely to
-# lead to a conflict than something that's human-readable (like "--streamlit-fake-field"
-# or something).
-PROTECTION_SUFFIX = "--p5bJXXpQgvPz6yvQMFiy"
-SEPARATED_INDEX_COLUMN_NAME = SEPARATED_INDEX_COLUMN_TITLE + PROTECTION_SUFFIX
-MELTED_Y_COLUMN_NAME = MELTED_Y_COLUMN_TITLE + PROTECTION_SUFFIX
-MELTED_COLOR_COLUMN_NAME = MELTED_COLOR_COLUMN_TITLE + PROTECTION_SUFFIX
-
-# Name we use for a column we know doesn't exist in the data, to address a Vega-Lite rendering bug
-# where empty charts need x, y encodings set in order to take up space.
-NON_EXISTENT_COLUMN_NAME = "DOES_NOT_EXIST" + PROTECTION_SUFFIX
-
-
-class ArrowAltairMixin:
     @gather_metrics("line_chart")
     def line_chart(
         self,
@@ -232,8 +341,8 @@ class ArrowAltairMixin:
            height: 440px
 
         """
-        proto = ArrowVegaLiteChartProto()
-        chart, add_rows_metadata = _generate_chart(
+
+        chart, add_rows_metadata = generate_chart(
             chart_type=ChartType.LINE,
             data=data,
             x_from_user=x,
@@ -243,10 +352,11 @@ class ArrowAltairMixin:
             width=width,
             height=height,
         )
-        marshall(proto, chart, use_container_width, theme="streamlit")
-
-        return self.dg._enqueue(
-            "arrow_line_chart", proto, add_rows_metadata=add_rows_metadata
+        return self._altair_chart(
+            chart,
+            use_container_width=use_container_width,
+            theme="streamlit",
+            add_rows_metadata=add_rows_metadata,
         )
 
     @gather_metrics("area_chart")
@@ -391,8 +501,7 @@ class ArrowAltairMixin:
 
         """
 
-        proto = ArrowVegaLiteChartProto()
-        chart, add_rows_metadata = _generate_chart(
+        chart, add_rows_metadata = generate_chart(
             chart_type=ChartType.AREA,
             data=data,
             x_from_user=x,
@@ -402,10 +511,11 @@ class ArrowAltairMixin:
             width=width,
             height=height,
         )
-        marshall(proto, chart, use_container_width, theme="streamlit")
-
-        return self.dg._enqueue(
-            "arrow_area_chart", proto, add_rows_metadata=add_rows_metadata
+        return self._altair_chart(
+            chart,
+            use_container_width=use_container_width,
+            theme="streamlit",
+            add_rows_metadata=add_rows_metadata,
         )
 
     @gather_metrics("bar_chart")
@@ -552,8 +662,7 @@ class ArrowAltairMixin:
 
         """
 
-        proto = ArrowVegaLiteChartProto()
-        chart, add_rows_metadata = _generate_chart(
+        chart, add_rows_metadata = generate_chart(
             chart_type=ChartType.BAR,
             data=data,
             x_from_user=x,
@@ -563,10 +672,11 @@ class ArrowAltairMixin:
             width=width,
             height=height,
         )
-        marshall(proto, chart, use_container_width, theme="streamlit")
-
-        return self.dg._enqueue(
-            "arrow_bar_chart", proto, add_rows_metadata=add_rows_metadata
+        return self._altair_chart(
+            chart,
+            use_container_width=use_container_width,
+            theme="streamlit",
+            add_rows_metadata=add_rows_metadata,
         )
 
     @gather_metrics("scatter_chart")
@@ -725,8 +835,8 @@ class ArrowAltairMixin:
            height: 440px
 
         """
-        proto = ArrowVegaLiteChartProto()
-        chart, add_rows_metadata = _generate_chart(
+
+        chart, add_rows_metadata = generate_chart(
             chart_type=ChartType.SCATTER,
             data=data,
             x_from_user=x,
@@ -736,10 +846,11 @@ class ArrowAltairMixin:
             width=width,
             height=height,
         )
-        marshall(proto, chart, use_container_width, theme="streamlit")
-
-        return self.dg._enqueue(
-            "arrow_scatter_chart", proto, add_rows_metadata=add_rows_metadata
+        return self._altair_chart(
+            chart,
+            use_container_width=use_container_width,
+            theme="streamlit",
+            add_rows_metadata=add_rows_metadata,
         )
 
     @gather_metrics("altair_chart")
@@ -790,780 +901,134 @@ class ArrowAltairMixin:
         https://altair-viz.github.io/gallery/.
 
         """
-        if theme != "streamlit" and theme != None:
-            raise StreamlitAPIException(
-                f'You set theme="{theme}" while Streamlit charts only support theme=”streamlit” or theme=None to fallback to the default library theme.'
-            )
-        proto = ArrowVegaLiteChartProto()
-        marshall(
-            proto,
-            altair_chart,
-            use_container_width=use_container_width,
-            theme=theme,
-        )
-
-        return self.dg._enqueue("arrow_vega_lite_chart", proto)
-
-    @property
-    def dg(self) -> DeltaGenerator:
-        """Get our DeltaGenerator."""
-        return cast("DeltaGenerator", self)
-
-
-def _is_date_column(df: pd.DataFrame, name: str | None) -> bool:
-    """True if the column with the given name stores datetime.date values.
-
-    This function just checks the first value in the given column, so
-    it's meaningful only for columns whose values all share the same type.
-
-    Parameters
-    ----------
-    df : pd.DataFrame
-    name : str
-        The column name
-
-    Returns
-    -------
-    bool
-
-    """
-    if name is None:
-        return False
-
-    column = df[name]
-    if column.size == 0:
-        return False
-
-    return isinstance(column.iloc[0], date)
-
-
-def _melt_data(
-    df: pd.DataFrame,
-    columns_to_leave_alone: list[str],
-    columns_to_melt: list[str] | None,
-    new_y_column_name: str,
-    new_color_column_name: str,
-) -> pd.DataFrame:
-    """Converts a wide-format dataframe to a long-format dataframe."""
-    import pandas as pd
-    from pandas.api.types import infer_dtype
-
-    melted_df = pd.melt(
-        df,
-        id_vars=columns_to_leave_alone,
-        value_vars=columns_to_melt,
-        var_name=new_color_column_name,
-        value_name=new_y_column_name,
-    )
-
-    y_series = melted_df[new_y_column_name]
-    if (
-        y_series.dtype == "object"
-        and "mixed" in infer_dtype(y_series)
-        and len(y_series.unique()) > 100
-    ):
-        raise StreamlitAPIException(
-            "The columns used for rendering the chart contain too many values with mixed types. Please select the columns manually via the y parameter."
-        )
-
-    # Arrow has problems with object types after melting two different dtypes
-    # pyarrow.lib.ArrowTypeError: "Expected a <TYPE> object, got a object"
-    fixed_df = type_util.fix_arrow_incompatible_column_types(
-        melted_df,
-        selected_columns=[
-            *columns_to_leave_alone,
-            new_color_column_name,
-            new_y_column_name,
-        ],
-    )
-
-    return fixed_df
-
-
-def prep_data(
-    df: pd.DataFrame,
-    x_column: str | None,
-    y_column_list: list[str],
-    color_column: str | None,
-    size_column: str | None,
-) -> tuple[pd.DataFrame, str | None, str | None, str | None, str | None]:
-    """Prepares the data for charting. This is also used in add_rows.
-
-    Returns the prepared dataframe and the new names of the x column (taking the index reset into
-    consideration) and y, color, and size columns.
-    """
-
-    # If y is provided, but x is not, we'll use the index as x.
-    # So we need to pull the index into its own column.
-    x_column = _maybe_reset_index_in_place(df, x_column, y_column_list)
-
-    # Drop columns we're not using.
-    selected_data = _drop_unused_columns(
-        df, x_column, color_column, size_column, *y_column_list
-    )
-
-    # Maybe convert color to Vega colors.
-    _maybe_convert_color_column_in_place(selected_data, color_column)
-
-    # Make sure all columns have string names.
-    (
-        x_column,
-        y_column_list,
-        color_column,
-        size_column,
-    ) = _convert_col_names_to_str_in_place(
-        selected_data, x_column, y_column_list, color_column, size_column
-    )
-
-    # Maybe melt data from wide format into long format.
-    melted_data, y_column, color_column = _maybe_melt(
-        selected_data, x_column, y_column_list, color_column, size_column
-    )
-
-    # Return the data, but also the new names to use for x, y, and color.
-    return melted_data, x_column, y_column, color_column, size_column
-
-
-def _generate_chart(
-    chart_type: ChartType,
-    data: Data | None,
-    x_from_user: str | None = None,
-    y_from_user: str | Sequence[str] | None = None,
-    color_from_user: str | Color | list[Color] | None = None,
-    size_from_user: str | float | None = None,
-    width: int = 0,
-    height: int = 0,
-) -> tuple[alt.Chart, AddRowsMetadata]:
-    """Function to use the chart's type, data columns and indices to figure out the chart's spec."""
-    import altair as alt
-
-    df = type_util.convert_anything_to_df(data, ensure_copy=True)
-
-    # From now on, use "df" instead of "data". Deleting "data" to guarantee we follow this.
-    del data
-
-    # Convert arguments received from the user to things Vega-Lite understands.
-    # Get name of column to use for x.
-    x_column = _parse_x_column(df, x_from_user)
-    # Get name of columns to use for y.
-    y_column_list = _parse_y_columns(df, y_from_user, x_column)
-    # Get name of column to use for color, or constant value to use. Any/both could be None.
-    color_column, color_value = _parse_generic_column(df, color_from_user)
-    # Get name of column to use for size, or constant value to use. Any/both could be None.
-    size_column, size_value = _parse_generic_column(df, size_from_user)
-
-    # Store some info so we can use it in add_rows.
-    add_rows_metadata = AddRowsMetadata(
-        # The last index of df so we can adjust the input df in add_rows:
-        last_index=last_index_for_melted_dataframes(df),
-        # This is the input to prep_data (except for the df):
-        columns=dict(
-            x_column=x_column,
-            y_column_list=y_column_list,
-            color_column=color_column,
-            size_column=size_column,
-        ),
-    )
-
-    # At this point, all foo_column variables are either None/empty or contain actual
-    # columns that are guaranteed to exist.
-
-    df, x_column, y_column, color_column, size_column = prep_data(
-        df, x_column, y_column_list, color_column, size_column
-    )
-
-    # At this point, x_column is only None if user did not provide one AND df is empty.
-
-    # Create a Chart with x and y encodings.
-    chart = alt.Chart(
-        data=df,
-        mark=chart_type.value["mark_type"],
-        width=width,
-        height=height,
-    ).encode(
-        x=_get_x_encoding(df, x_column, x_from_user, chart_type),
-        y=_get_y_encoding(df, y_column, y_from_user),
-    )
-
-    # Set up opacity encoding.
-    opacity_enc = _get_opacity_encoding(chart_type, color_column)
-    if opacity_enc is not None:
-        chart = chart.encode(opacity=opacity_enc)
-
-    # Set up color encoding.
-    color_enc = _get_color_encoding(
-        df, color_value, color_column, y_column_list, color_from_user
-    )
-    if color_enc is not None:
-        chart = chart.encode(color=color_enc)
-
-    # Set up size encoding.
-    size_enc = _get_size_encoding(chart_type, size_column, size_value)
-    if size_enc is not None:
-        chart = chart.encode(size=size_enc)
-
-    # Set up tooltip encoding.
-    if x_column is not None and y_column is not None:
-        chart = chart.encode(
-            tooltip=_get_tooltip_encoding(
-                x_column,
-                y_column,
-                size_column,
-                color_column,
-                color_enc,
-            )
-        )
-
-    return chart.interactive(), add_rows_metadata
-
-
-def _maybe_reset_index_in_place(
-    df: pd.DataFrame, x_column: str | None, y_column_list: list[str]
-) -> str | None:
-    if x_column is None and len(y_column_list) > 0:
-        if df.index.name is None:
-            # Pick column name that is unlikely to collide with user-given names.
-            x_column = SEPARATED_INDEX_COLUMN_NAME
-        else:
-            # Reuse index's name for the new column.
-            x_column = df.index.name
-
-        df.index.name = x_column
-        df.reset_index(inplace=True)
-
-    return x_column
-
-
-def _drop_unused_columns(df: pd.DataFrame, *column_names: str | None) -> pd.DataFrame:
-    """Returns a subset of df, selecting only column_names that aren't None."""
-
-    # We can't just call set(col_names) because sets don't have stable ordering,
-    # which means tests that depend on ordering will fail.
-    # Performance-wise, it's not a problem, though, since this function is only ever
-    # used on very small lists.
-    seen = set()
-    keep = []
-
-    for x in column_names:
-        if x is None:
-            continue
-        if x in seen:
-            continue
-        seen.add(x)
-        keep.append(x)
-
-    return df[keep]
-
-
-def _maybe_convert_color_column_in_place(df: pd.DataFrame, color_column: str | None):
-    """If needed, convert color column to a format Vega understands."""
-    if color_column is None or len(df[color_column]) == 0:
-        return
-
-    first_color_datum = df[color_column].iat[0]
-
-    if is_hex_color_like(first_color_datum):
-        # Hex is already CSS-valid.
-        pass
-    elif is_color_tuple_like(first_color_datum):
-        # Tuples need to be converted to CSS-valid.
-        df[color_column] = df[color_column].map(to_css_color)
-    else:
-        # Other kinds of colors columns (i.e. pure numbers or nominal strings) shouldn't
-        # be converted since they are treated by Vega-Lite as sequential or categorical colors.
-        pass
-
-
-def _convert_col_names_to_str_in_place(
-    df: pd.DataFrame,
-    x_column: str | None,
-    y_column_list: list[str],
-    color_column: str | None,
-    size_column: str | None,
-) -> tuple[str | None, list[str], str | None, str | None]:
-    """Converts column names to strings, since Vega-Lite does not accept ints, etc."""
-    import pandas as pd
-
-    column_names = list(df.columns)  # list() converts RangeIndex, etc, to regular list.
-    str_column_names = [str(c) for c in column_names]
-    df.columns = pd.Index(str_column_names)
-
-    return (
-        None if x_column is None else str(x_column),
-        [str(c) for c in y_column_list],
-        None if color_column is None else str(color_column),
-        None if size_column is None else str(size_column),
-    )
-
-
-def _parse_generic_column(
-    df: pd.DataFrame, column_or_value: Any
-) -> tuple[str | None, Any]:
-    if isinstance(column_or_value, str) and column_or_value in df.columns:
-        column_name = column_or_value
-        value = None
-    else:
-        column_name = None
-        value = column_or_value
-
-    return column_name, value
-
-
-def _parse_x_column(df: pd.DataFrame, x_from_user: str | None) -> str | None:
-    if x_from_user is None:
-        return None
-
-    elif isinstance(x_from_user, str):
-        if x_from_user not in df.columns:
-            raise StreamlitColumnNotFoundError(df, x_from_user)
-
-        return x_from_user
-
-    else:
-        raise StreamlitAPIException(
-            "x parameter should be a column name (str) or None to use the "
-            f" dataframe's index. Value given: {x_from_user} "
-            f"(type {type(x_from_user)})"
+        return self._altair_chart(
+            altair_chart, use_container_width=use_container_width, theme=theme
         )
 
+    @gather_metrics("vega_lite_chart")
+    def vega_lite_chart(
+        self,
+        data: Data = None,
+        spec: dict[str, Any] | None = None,
+        use_container_width: bool = False,
+        theme: Literal["streamlit"] | None = "streamlit",
+        **kwargs: Any,
+    ) -> DeltaGenerator:
+        """Display a chart using the Vega-Lite library.
 
-def _parse_y_columns(
-    df: pd.DataFrame,
-    y_from_user: str | Sequence[str] | None,
-    x_column: str | None,
-) -> list[str]:
-    y_column_list: list[str] = []
-
-    if y_from_user is None:
-        y_column_list = list(df.columns)
-
-    elif isinstance(y_from_user, str):
-        y_column_list = [y_from_user]
-
-    elif type_util.is_sequence(y_from_user):
-        y_column_list = list(str(col) for col in y_from_user)
-
-    else:
-        raise StreamlitAPIException(
-            "y parameter should be a column name (str) or list thereof. "
-            f"Value given: {y_from_user} (type {type(y_from_user)})"
-        )
-
-    for col in y_column_list:
-        if col not in df.columns:
-            raise StreamlitColumnNotFoundError(df, col)
+        Parameters
+        ----------
+        data : pandas.DataFrame, pandas.Styler, pyarrow.Table, numpy.ndarray, Iterable, dict, or None
+            Either the data to be plotted or a Vega-Lite spec containing the
+            data (which more closely follows the Vega-Lite API).
 
-    # y_column_list should only include x_column when user explicitly asked for it.
-    if x_column in y_column_list and (not y_from_user or x_column not in y_from_user):
-        y_column_list.remove(x_column)
+        spec : dict or None
+            The Vega-Lite spec for the chart. If the spec was already passed in
+            the previous argument, this must be set to None. See
+            https://vega.github.io/vega-lite/docs/ for more info.
 
-    return y_column_list
+        use_container_width : bool
+            If True, set the chart width to the column width. This takes
+            precedence over Vega-Lite's native `width` value.
 
+        theme : "streamlit" or None
+            The theme of the chart. Currently, we only support "streamlit" for the Streamlit
+            defined design or None to fallback to the default behavior of the library.
 
-def _get_opacity_encoding(
-    chart_type: ChartType, color_column: str | None
-) -> alt.OpacityValue | None:
-    import altair as alt
+        **kwargs : any
+            Same as spec, but as keywords.
 
-    if color_column and chart_type == ChartType.AREA:
-        return alt.OpacityValue(0.7)
+        Example
+        -------
+        >>> import streamlit as st
+        >>> import pandas as pd
+        >>> import numpy as np
+        >>>
+        >>> chart_data = pd.DataFrame(np.random.randn(200, 3), columns=["a", "b", "c"])
+        >>>
+        >>> st.vega_lite_chart(
+        ...    chart_data,
+        ...    {
+        ...        "mark": {"type": "circle", "tooltip": True},
+        ...        "encoding": {
+        ...            "x": {"field": "a", "type": "quantitative"},
+        ...            "y": {"field": "b", "type": "quantitative"},
+        ...            "size": {"field": "c", "type": "quantitative"},
+        ...            "color": {"field": "c", "type": "quantitative"},
+        ...        },
+        ...    },
+        ... )
 
-    return None
+        .. output::
+           https://doc-vega-lite-chart.streamlit.app/
+           height: 300px
 
+        Examples of Vega-Lite usage without Streamlit can be found at
+        https://vega.github.io/vega-lite/examples/. Most of those can be easily
+        translated to the syntax shown above.
 
-def _get_axis_config(df: pd.DataFrame, column_name: str | None, grid: bool) -> alt.Axis:
-    import altair as alt
-    from pandas.api.types import is_integer_dtype
-
-    if column_name is not None and is_integer_dtype(df[column_name]):
-        # Use a max tick size of 1 for integer columns (prevents zoom into float numbers)
-        # and deactivate grid lines for x-axis
-        return alt.Axis(tickMinStep=1, grid=grid)
-
-    return alt.Axis(grid=grid)
-
-
-def _maybe_melt(
-    df: pd.DataFrame,
-    x_column: str | None,
-    y_column_list: list[str],
-    color_column: str | None,
-    size_column: str | None,
-) -> tuple[pd.DataFrame, str | None, str | None]:
-    """If multiple columns are set for y, melt the dataframe into long format."""
-    y_column: str | None
-
-    if len(y_column_list) == 0:
-        y_column = None
-    elif len(y_column_list) == 1:
-        y_column = y_column_list[0]
-    elif x_column is not None:
-        # Pick column names that are unlikely to collide with user-given names.
-        y_column = MELTED_Y_COLUMN_NAME
-        color_column = MELTED_COLOR_COLUMN_NAME
-
-        columns_to_leave_alone = [x_column]
-        if size_column:
-            columns_to_leave_alone.append(size_column)
-
-        df = _melt_data(
-            df=df,
-            columns_to_leave_alone=columns_to_leave_alone,
-            columns_to_melt=y_column_list,
-            new_y_column_name=y_column,
-            new_color_column_name=color_column,
+        """
+        return self._vega_lite_chart(
+            data=data,
+            spec=spec,
+            use_container_width=use_container_width,
+            theme=theme,
+            **kwargs,
         )
 
-    return df, y_column, color_column
-
-
-def _get_x_encoding(
-    df: pd.DataFrame,
-    x_column: str | None,
-    x_from_user: str | None,
-    chart_type: ChartType,
-) -> alt.X:
-    import altair as alt
-
-    if x_column is None:
-        # If no field is specified, the full axis disappears when no data is present.
-        # Maybe a bug in vega-lite? So we pass a field that doesn't exist.
-        x_field = NON_EXISTENT_COLUMN_NAME
-        x_title = ""
-    elif x_column == SEPARATED_INDEX_COLUMN_NAME:
-        # If the x column name is the crazy anti-collision name we gave it, then need to set
-        # up a title so we never show the crazy name to the user.
-        x_field = x_column
-        # Don't show a label in the x axis (not even a nice label like
-        # SEPARATED_INDEX_COLUMN_TITLE) when we pull the x axis from the index.
-        x_title = ""
-    else:
-        x_field = x_column
-
-        # Only show a label in the x axis if the user passed a column explicitly. We
-        # could go either way here, but I'm keeping this to avoid breaking the existing
-        # behavior.
-        if x_from_user is None:
-            x_title = ""
-        else:
-            x_title = x_column
-
-    return alt.X(
-        x_field,
-        title=x_title,
-        type=_get_x_encoding_type(df, chart_type, x_column),
-        scale=alt.Scale(),
-        axis=_get_axis_config(df, x_column, grid=False),
-    )
-
-
-def _get_y_encoding(
-    df: pd.DataFrame,
-    y_column: str | None,
-    y_from_user: str | Sequence[str] | None,
-) -> alt.Y:
-    import altair as alt
-
-    if y_column is None:
-        # If no field is specified, the full axis disappears when no data is present.
-        # Maybe a bug in vega-lite? So we pass a field that doesn't exist.
-        y_field = NON_EXISTENT_COLUMN_NAME
-        y_title = ""
-    elif y_column == MELTED_Y_COLUMN_NAME:
-        # If the y column name is the crazy anti-collision name we gave it, then need to set
-        # up a title so we never show the crazy name to the user.
-        y_field = y_column
-        # Don't show a label in the y axis (not even a nice label like
-        # MELTED_Y_COLUMN_TITLE) when we pull the x axis from the index.
-        y_title = ""
-    else:
-        y_field = y_column
-
-        # Only show a label in the y axis if the user passed a column explicitly. We
-        # could go either way here, but I'm keeping this to avoid breaking the existing
-        # behavior.
-        if y_from_user is None:
-            y_title = ""
-        else:
-            y_title = y_column
-
-    return alt.Y(
-        field=y_field,
-        title=y_title,
-        type=_get_y_encoding_type(df, y_column),
-        scale=alt.Scale(),
-        axis=_get_axis_config(df, y_column, grid=True),
-    )
-
-
-def _get_color_encoding(
-    df: pd.DataFrame,
-    color_value: Color | None,
-    color_column: str | None,
-    y_column_list: list[str],
-    color_from_user: str | Color | list[Color] | None,
-) -> alt.Color | alt.ColorValue | None:
-    import altair as alt
-
-    has_color_value = color_value not in [None, [], tuple()]
-
-    # If user passed a color value, that should win over colors coming from the
-    # color column (be they manual or auto-assigned due to melting)
-    if has_color_value:
-        # If the color value is color-like, return that.
-        if is_color_like(cast(Any, color_value)):
-            if len(y_column_list) != 1:
-                raise StreamlitColorLengthError([color_value], y_column_list)
-
-            return alt.ColorValue(to_css_color(cast(Any, color_value)))
-
-        # If the color value is a list of colors of approriate length, return that.
-        elif isinstance(color_value, (list, tuple)):
-            color_values = cast(Collection[Color], color_value)
-
-            if len(color_values) != len(y_column_list):
-                raise StreamlitColorLengthError(color_values, y_column_list)
-
-            if len(color_value) == 1:
-                return alt.ColorValue(to_css_color(cast(Any, color_value[0])))
-            else:
-                return alt.Color(
-                    field=color_column,
-                    scale=alt.Scale(range=[to_css_color(c) for c in color_values]),
-                    legend=COLOR_LEGEND_SETTINGS,
-                    type="nominal",
-                    title=" ",
-                )
-
-        raise StreamlitInvalidColorError(df, color_from_user)
-
-    elif color_column is not None:
-        column_type: str | tuple[str, list[Any]]
-
-        if color_column == MELTED_COLOR_COLUMN_NAME:
-            column_type = "nominal"
-        else:
-            column_type = type_util.infer_vegalite_type(df[color_column])
-
-        color_enc = alt.Color(
-            field=color_column, legend=COLOR_LEGEND_SETTINGS, type=column_type
+    def _altair_chart(
+        self,
+        altair_chart: alt.Chart,
+        use_container_width: bool = False,
+        theme: Literal["streamlit"] | None = "streamlit",
+        add_rows_metadata: AddRowsMetadata | None = None,
+    ) -> DeltaGenerator:
+        """Internal method to enqueue a vega-lite chart element based on an Altair chart."""
+        vega_lite_spec = _convert_altair_to_vega_lite_spec(altair_chart)
+        return self._vega_lite_chart(
+            data=None,  # The data is already part of the spec
+            spec=vega_lite_spec,
+            use_container_width=use_container_width,
+            theme=theme,
+            add_rows_metadata=add_rows_metadata,
         )
 
-        # Fix title if DF was melted
-        if color_column == MELTED_COLOR_COLUMN_NAME:
-            # This has to contain an empty space, otherwise the
-            # full y-axis disappears (maybe a bug in vega-lite)?
-            color_enc["title"] = " "
-
-        # If the 0th element in the color column looks like a color, we'll use the color column's
-        # values as the colors in our chart.
-        elif len(df[color_column]) and is_color_like(df[color_column].iat[0]):
-            color_range = [to_css_color(c) for c in df[color_column].unique()]
-            color_enc["scale"] = alt.Scale(range=color_range)
-            # Don't show the color legend, because it will just show text with the color values,
-            # like #f00, #00f, etc, which are not user-readable.
-            color_enc["legend"] = None
-
-        # Otherwise, let Vega-Lite auto-assign colors.
-        # This codepath is typically reached when the color column contains numbers (in which case
-        # Vega-Lite uses a color gradient to represent them) or strings (in which case Vega-Lite
-        # assigns one color for each unique value).
-        else:
-            pass
-
-        return color_enc
-
-    return None
-
-
-def _get_size_encoding(
-    chart_type: ChartType,
-    size_column: str | None,
-    size_value: str | float | None,
-) -> alt.Size | alt.SizeValue | None:
-    import altair as alt
-
-    if chart_type == ChartType.SCATTER:
-        if size_column is not None:
-            return alt.Size(
-                size_column,
-                legend=SIZE_LEGEND_SETTINGS,
-            )
+    def _vega_lite_chart(
+        self,
+        data: Data = None,
+        spec: dict[str, Any] | None = None,
+        use_container_width: bool = False,
+        theme: Literal["streamlit"] | None = "streamlit",
+        add_rows_metadata: AddRowsMetadata | None = None,
+        **kwargs: Any,
+    ) -> DeltaGenerator:
+        """Internal method to enqueue a vega-lite chart element based on a vega-lite spec."""
 
-        elif isinstance(size_value, (float, int)):
-            return alt.SizeValue(size_value)
-        elif size_value is None:
-            return alt.SizeValue(100)
-        else:
+        if theme not in ["streamlit", None]:
             raise StreamlitAPIException(
-                f"This does not look like a valid size: {repr(size_value)}"
-            )
-
-    elif size_column is not None or size_value is not None:
-        raise Error(
-            f"Chart type {chart_type.name} does not support size argument. "
-            "This should never happen!"
-        )
-
-    return None
-
-
-def _get_tooltip_encoding(
-    x_column: str,
-    y_column: str,
-    size_column: str | None,
-    color_column: str | None,
-    color_enc: alt.Color | alt.ColorValue | None,
-) -> list[alt.Tooltip]:
-    import altair as alt
-
-    tooltip = []
-
-    # If the x column name is the crazy anti-collision name we gave it, then need to set
-    # up a tooltip title so we never show the crazy name to the user.
-    if x_column == SEPARATED_INDEX_COLUMN_NAME:
-        tooltip.append(alt.Tooltip(x_column, title=SEPARATED_INDEX_COLUMN_TITLE))
-    else:
-        tooltip.append(alt.Tooltip(x_column))
-
-    # If the y column name is the crazy anti-collision name we gave it, then need to set
-    # up a tooltip title so we never show the crazy name to the user.
-    if y_column == MELTED_Y_COLUMN_NAME:
-        tooltip.append(
-            alt.Tooltip(
-                y_column,
-                title=MELTED_Y_COLUMN_TITLE,
-                type="quantitative",  # Just picked something random. Doesn't really matter!
-            )
-        )
-    else:
-        tooltip.append(alt.Tooltip(y_column))
-
-    # If we earlier decided that there should be no color legend, that's because the
-    # user passed a color column with actual color values (like "#ff0"), so we should
-    # not show the color values in the tooltip.
-    if color_column and getattr(color_enc, "legend", True) is not None:
-        # Use a human-readable title for the color.
-        if color_column == MELTED_COLOR_COLUMN_NAME:
-            tooltip.append(
-                alt.Tooltip(
-                    color_column,
-                    title=MELTED_COLOR_COLUMN_TITLE,
-                    type="nominal",
-                )
+                f'You set theme="{theme}" while Streamlit charts only support theme=”streamlit” or theme=None to fallback to the default library theme.'
             )
-        else:
-            tooltip.append(alt.Tooltip(color_column))
-
-    if size_column:
-        tooltip.append(alt.Tooltip(size_column))
-
-    return tooltip
-
 
-def _get_x_encoding_type(
-    df: pd.DataFrame, chart_type: ChartType, x_column: str | None
-) -> type_util.VegaLiteType:
-    if x_column is None:
-        return "quantitative"  # Anything. If None, Vega-Lite may hide the axis.
+        # Support passing data inside spec['datasets'] and spec['data'].
+        # (The data gets pulled out of the spec dict later on.)
+        if isinstance(data, dict) and spec is None:
+            spec = data
+            data = None
 
-    # Bar charts should have a discrete (ordinal) x-axis, UNLESS type is date/time
-    # https://github.com/streamlit/streamlit/pull/2097#issuecomment-714802475
-    if chart_type == ChartType.BAR and not _is_date_column(df, x_column):
-        return "ordinal"
-
-    return type_util.infer_vegalite_type(df[x_column])
-
-
-def _get_y_encoding_type(
-    df: pd.DataFrame, y_column: str | None
-) -> type_util.VegaLiteType:
-    if y_column:
-        return type_util.infer_vegalite_type(df[y_column])
-
-    return "quantitative"  # Pick anything. If undefined, Vega-Lite may hide the axis.
-
-
-def marshall(
-    vega_lite_chart: ArrowVegaLiteChartProto,
-    altair_chart: alt.Chart,
-    use_container_width: bool = False,
-    theme: None | Literal["streamlit"] = "streamlit",
-    **kwargs: Any,
-) -> None:
-    """Marshall chart's data into proto."""
-    import altair as alt
-
-    # Normally altair_chart.to_dict() would transform the dataframe used by the
-    # chart into an array of dictionaries. To avoid that, we install a
-    # transformer that replaces datasets with a reference by the object id of
-    # the dataframe. We then fill in the dataset manually later on.
-
-    datasets = {}
-
-    def id_transform(data) -> dict[str, str]:
-        """Altair data transformer that returns a fake named dataset with the
-        object id.
-        """
-        name = str(id(data))
-        datasets[name] = data
-        return {"name": name}
-
-    alt.data_transformers.register("id", id_transform)  # type: ignore[attr-defined,unused-ignore]
-
-    # The default altair theme has some width/height defaults defined
-    # which are not useful for Streamlit. Therefore, we change the theme to
-    # "none" to avoid those defaults.
-    with alt.themes.enable("none") if alt.themes.active == "default" else nullcontext():  # type: ignore[attr-defined,unused-ignore]
-        with alt.data_transformers.enable("id"):  # type: ignore[attr-defined,unused-ignore]
-            chart_dict = altair_chart.to_dict()
-
-            # Put datasets back into the chart dict but note how they weren't
-            # transformed.
-            chart_dict["datasets"] = datasets
+        proto = ArrowVegaLiteChartProto()
 
-            arrow_vega_lite.marshall(
-                vega_lite_chart,
-                chart_dict,
-                use_container_width=use_container_width,
-                theme=theme,
-                **kwargs,
-            )
+        spec = _prepare_vega_lite_spec(spec, use_container_width, **kwargs)
+        _marshall_chart_data(proto, spec, data)
 
+        proto.spec = json.dumps(spec)
+        proto.use_container_width = use_container_width
+        proto.theme = theme or ""
 
-class StreamlitColumnNotFoundError(StreamlitAPIException):
-    def __init__(self, df, col_name, *args):
-        available_columns = ", ".join(str(c) for c in list(df.columns))
-        message = (
-            f'Data does not have a column named `"{col_name}"`. '
-            f"Available columns are `{available_columns}`"
+        return self.dg._enqueue(
+            "arrow_vega_lite_chart", proto, add_rows_metadata=add_rows_metadata
         )
-        super().__init__(message, *args)
-
-
-class StreamlitInvalidColorError(StreamlitAPIException):
-    def __init__(self, df, color_from_user, *args):
-        ", ".join(str(c) for c in list(df.columns))
-        message = f"""
-This does not look like a valid color argument: `{color_from_user}`.
-
-The color argument can be:
 
-* A hex string like "#ffaa00" or "#ffaa0088".
-* An RGB or RGBA tuple with the red, green, blue, and alpha
-  components specified as ints from 0 to 255 or floats from 0.0 to
-  1.0.
-* The name of a column.
-* Or a list of colors, matching the number of y columns to draw.
-        """
-        super().__init__(message, *args)
-
-
-class StreamlitColorLengthError(StreamlitAPIException):
-    def __init__(self, color_values, y_column_list, *args):
-        message = (
-            f"The list of colors `{color_values}` must have the same "
-            "length as the list of columns to be colored "
-            f"`{y_column_list}`."
-        )
-        super().__init__(message, *args)
+    @property
+    def dg(self) -> DeltaGenerator:
+        """Get our DeltaGenerator."""
+        return cast("DeltaGenerator", self)