streamlit-nightly 1.37.2.dev20240810__py2.py3-none-any.whl → 1.37.2.dev20240813__py2.py3-none-any.whl

streamlit/connections/sql_connection.py

@@ -12,6 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# NOTE: We ignore all mypy import-not-found errors as top-level since
+# this module is optional and the SQLAlchemy dependency is not installed
+# by default.
+# mypy: disable-error-code="import-not-found, redundant-cast"
+
 from __future__ import annotations
 
 from collections import ChainMap
@@ -277,7 +282,7 @@ class SQLConnection(BaseConnection["Engine"]):
 
         This is equivalent to accessing ``self._instance.driver``.
         """
-        return self._instance.driver
+        return cast(str, self._instance.driver)
 
     @property
     def session(self) -> Session:

streamlit/dataframe_util.py

@@ -42,6 +42,7 @@ from typing_extensions import TypeAlias, TypeGuard
 
 from streamlit import config, errors, logger, string_util
 from streamlit.type_util import (
+    NumpyShape,
     has_callable_attr,
     is_custom_dict,
     is_dataclass_instance,
@@ -81,6 +82,11 @@ _SNOWPANDAS_INDEX_TYPE_STR: Final = (
 _POLARS_DATAFRAME: Final = "polars.dataframe.frame.DataFrame"
 _POLARS_SERIES: Final = "polars.series.series.Series"
 _POLARS_LAZYFRAME: Final = "polars.lazyframe.frame.LazyFrame"
+_DASK_DATAFRAME: Final = "dask.dataframe.core.DataFrame"
+_DASK_SERIES: Final = "dask.dataframe.core.Series"
+_DASK_INDEX: Final = "dask.dataframe.core.Index"
+_RAY_MATERIALIZED_DATASET: Final = "ray.data.dataset.MaterializedDataset"
+_RAY_DATASET: Final = "ray.data.dataset.Dataset"
 
 V_co = TypeVar(
     "V_co",
@@ -116,7 +122,7 @@ Data: TypeAlias = Union[
     "Styler",
     "Index",
     "pa.Table",
-    "np.ndarray",
+    "np.ndarray[Any, np.dtype[Any]]",
     Iterable[Any],
     Dict[Any, Any],
     None,
@@ -146,6 +152,8 @@ class DataFormat(Enum):
     POLARS_SERIES = auto()  # polars.series.series.Series
     XARRAY_DATASET = auto()  # xarray.Dataset
     XARRAY_DATA_ARRAY = auto()  # xarray.DataArray
+    DASK_OBJECT = auto()  # dask.dataframe.core.DataFrame, Series, Index
+    RAY_DATASET = auto()  # ray.data.dataset.Dataset, MaterializedDataset
     LIST_OF_RECORDS = auto()  # List[Dict[str, Scalar]]
     LIST_OF_ROWS = auto()  # List[List[Scalar]]
     LIST_OF_VALUES = auto()  # List[Scalar]
@@ -189,6 +197,8 @@ def is_dataframe_like(obj: object) -> bool:
         DataFormat.POLARS_LAZYFRAME,
         DataFormat.XARRAY_DATASET,
         DataFormat.XARRAY_DATA_ARRAY,
+        DataFormat.DASK_OBJECT,
+        DataFormat.RAY_DATASET,
         DataFormat.COLUMN_SERIES_MAPPING,
     ]
 
@@ -201,6 +211,8 @@ def is_unevaluated_data_object(obj: object) -> bool:
     - PySpark DataFrame
     - Modin DataFrame / Series
     - Snowpandas DataFrame / Series / Index
+    - Dask DataFrame / Series / Index
+    - Ray Dataset
     - Polars LazyFrame
     - Generator functions
 
@@ -213,7 +225,9 @@ def is_unevaluated_data_object(obj: object) -> bool:
         or is_pyspark_data_object(obj)
         or is_snowpandas_data_object(obj)
         or is_modin_data_object(obj)
+        or is_ray_dataset(obj)
         or is_polars_lazyframe(obj)
+        or is_dask_object(obj)
         or inspect.isgeneratorfunction(obj)
     )
 
@@ -247,6 +261,15 @@ def is_pyspark_data_object(obj: object) -> bool:
     )
 
 
+def is_dask_object(obj: object) -> bool:
+    """True if obj is a Dask DataFrame, Series, or Index."""
+    return (
+        is_type(obj, _DASK_DATAFRAME)
+        or is_type(obj, _DASK_SERIES)
+        or is_type(obj, _DASK_INDEX)
+    )
+
+
 def is_modin_data_object(obj: object) -> bool:
     """True if obj is of Modin Dataframe or Series"""
     return is_type(obj, _MODIN_DF_TYPE_STR) or is_type(obj, _MODIN_SERIES_TYPE_STR)
@@ -286,6 +309,11 @@ def is_polars_lazyframe(obj: object) -> bool:
     return is_type(obj, _POLARS_LAZYFRAME)
 
 
+def is_ray_dataset(obj: object) -> bool:
+    """True if obj is a Ray Dataset."""
+    return is_type(obj, _RAY_DATASET) or is_type(obj, _RAY_MATERIALIZED_DATASET)
+
+
 def is_pandas_styler(obj: object) -> TypeGuard[Styler]:
     """True if obj is a pandas Styler."""
     return is_type(obj, _PANDAS_STYLER_TYPE_STR)
@@ -436,6 +464,31 @@ def convert_anything_to_pandas_df(
             data = data.copy(deep=True)
         return data.to_series().to_frame()
 
+    if is_dask_object(data):
+        data = data.head(max_unevaluated_rows, compute=True)
+
+        # Dask returns a Pandas object (DataFrame, Series, Index) when
+        # executing operations like `head`.
+        if isinstance(data, (pd.Series, pd.Index)):
+            data = data.to_frame()
+
+        if data.shape[0] == max_unevaluated_rows:
+            _show_data_information(
+                f"⚠️ Showing only {string_util.simplify_number(max_unevaluated_rows)} "
+                "rows. Call `compute()` on the data object to show more."
+            )
+        return cast(pd.DataFrame, data)
+
+    if is_ray_dataset(data):
+        data = data.limit(max_unevaluated_rows).to_pandas()
+
+        if data.shape[0] == max_unevaluated_rows:
+            _show_data_information(
+                f"⚠️ Showing only {string_util.simplify_number(max_unevaluated_rows)} "
+                "rows. Call `to_pandas()` on the dataset to show more."
+            )
+        return cast(pd.DataFrame, data)
+
     if is_modin_data_object(data):
         data = data.head(max_unevaluated_rows)._to_pandas()
 
@@ -952,7 +1005,7 @@ def determine_data_format(input_data: Any) -> DataFormat:
     elif isinstance(input_data, pd.DataFrame):
         return DataFormat.PANDAS_DATAFRAME
     elif isinstance(input_data, np.ndarray):
-        if len(input_data.shape) == 1:
+        if len(cast(NumpyShape, input_data.shape)) == 1:
             # For technical reasons, we need to distinguish one
             # one-dimensional numpy array from multidimensional ones.
             return DataFormat.NUMPY_LIST
@@ -985,6 +1038,10 @@ def determine_data_format(input_data: Any) -> DataFormat:
         return DataFormat.XARRAY_DATASET
     elif is_xarray_data_array(input_data):
         return DataFormat.XARRAY_DATA_ARRAY
+    elif is_ray_dataset(input_data):
+        return DataFormat.RAY_DATASET
+    elif is_dask_object(input_data):
+        return DataFormat.DASK_OBJECT
     elif is_snowpark_data_object(input_data) or is_snowpark_row_list(input_data):
         return DataFormat.SNOWPARK_OBJECT
     elif (
@@ -1105,6 +1162,8 @@ def convert_pandas_df_to_data_format(
         DataFormat.PANDAS_ARRAY,
         DataFormat.MODIN_OBJECT,
         DataFormat.SNOWPANDAS_OBJECT,
+        DataFormat.DASK_OBJECT,
+        DataFormat.RAY_DATASET,
     ]:
         return df
     elif data_format == DataFormat.NUMPY_LIST:
@@ -1133,7 +1192,7 @@ def convert_pandas_df_to_data_format(
         data_format == DataFormat.POLARS_DATAFRAME
         or data_format == DataFormat.POLARS_LAZYFRAME
     ):
-        import polars as pl
+        import polars as pl  # type: ignore[import-not-found]
 
         return pl.from_pandas(df)
     elif data_format == DataFormat.POLARS_SERIES:
@@ -1141,7 +1200,7 @@ def convert_pandas_df_to_data_format(
 
         return pl.from_pandas(_pandas_df_to_series(df))
     elif data_format == DataFormat.XARRAY_DATASET:
-        import xarray as xr
+        import xarray as xr  # type: ignore[import-not-found]
 
         return xr.Dataset.from_dataframe(df)
     elif data_format == DataFormat.XARRAY_DATA_ARRAY:

streamlit/delta_generator.py

@@ -24,7 +24,6 @@ from typing import (
     Any,
     Callable,
     Final,
-    Hashable,
     Iterable,
     Literal,
     NoReturn,
@@ -38,7 +37,6 @@ from streamlit import (
     cli_util,
     config,
     cursor,
-    dataframe_util,
     env_util,
     logger,
     runtime,
@@ -89,26 +87,23 @@ from streamlit.elements.widgets.slider import SliderMixin
 from streamlit.elements.widgets.text_widgets import TextWidgetsMixin
 from streamlit.elements.widgets.time_widgets import TimeWidgetsMixin
 from streamlit.elements.write import WriteMixin
-from streamlit.errors import NoSessionContext, StreamlitAPIException
+from streamlit.errors import StreamlitAPIException
 from streamlit.proto import Block_pb2, ForwardMsg_pb2
 from streamlit.proto.RootContainer_pb2 import RootContainer
 from streamlit.runtime import caching
+from streamlit.runtime.scriptrunner import enqueue_message as _enqueue_message
 from streamlit.runtime.scriptrunner import get_script_run_ctx
 
 if TYPE_CHECKING:
     from google.protobuf.message import Message
-    from numpy import typing as npt
-    from pandas import DataFrame
 
     from streamlit.cursor import Cursor
-    from streamlit.dataframe_util import Data
     from streamlit.elements.lib.built_in_chart_utils import AddRowsMetadata
 
 
 MAX_DELTA_BYTES: Final[int] = 14 * 1024 * 1024  # 14MB
 
 Value = TypeVar("Value")
-DG = TypeVar("DG", bound="DeltaGenerator")
 
 # Type aliases for Ancestor Block Types
 BlockType: TypeAlias = str
@@ -554,121 +549,6 @@ class DeltaGenerator(
 
         return block_dg
 
-    def _arrow_add_rows(
-        self: DG,
-        data: Data = None,
-        **kwargs: (
-            DataFrame | npt.NDArray[Any] | Iterable[Any] | dict[Hashable, Any] | None
-        ),
-    ) -> DG | None:
-        """Concatenate a dataframe to the bottom of the current one.
-
-        Parameters
-        ----------
-        data : pandas.DataFrame, pandas.Styler, numpy.ndarray, Iterable, dict, or None
-            Table to concat. Optional.
-
-        **kwargs : pandas.DataFrame, numpy.ndarray, Iterable, dict, or None
-            The named dataset to concat. Optional. You can only pass in 1
-            dataset (including the one in the data parameter).
-
-        Example
-        -------
-        >>> import streamlit as st
-        >>> import pandas as pd
-        >>> import numpy as np
-        >>>
-        >>> df1 = pd.DataFrame(
-        ...     np.random.randn(50, 20), columns=("col %d" % i for i in range(20))
-        ... )
-        >>> my_table = st.table(df1)
-        >>>
-        >>> df2 = pd.DataFrame(
-        ...     np.random.randn(50, 20), columns=("col %d" % i for i in range(20))
-        ... )
-        >>> my_table.add_rows(df2)
-        >>> # Now the table shown in the Streamlit app contains the data for
-        >>> # df1 followed by the data for df2.
-
-        You can do the same thing with plots. For example, if you want to add
-        more data to a line chart:
-
-        >>> # Assuming df1 and df2 from the example above still exist...
-        >>> my_chart = st.line_chart(df1)
-        >>> my_chart.add_rows(df2)
-        >>> # Now the chart shown in the Streamlit app contains the data for
-        >>> # df1 followed by the data for df2.
-
-        And for plots whose datasets are named, you can pass the data with a
-        keyword argument where the key is the name:
-
-        >>> my_chart = st.vega_lite_chart(
-        ...     {
-        ...         "mark": "line",
-        ...         "encoding": {"x": "a", "y": "b"},
-        ...         "datasets": {
-        ...             "some_fancy_name": df1,  # <-- named dataset
-        ...         },
-        ...         "data": {"name": "some_fancy_name"},
-        ...     }
-        ... )
-        >>> my_chart.add_rows(some_fancy_name=df2)  # <-- name used as keyword
-
-        """
-        if self._root_container is None or self._cursor is None:
-            return self
-
-        if not self._cursor.is_locked:
-            raise StreamlitAPIException("Only existing elements can `add_rows`.")
-
-        # Accept syntax st._arrow_add_rows(df).
-        if data is not None and len(kwargs) == 0:
-            name = ""
-        # Accept syntax st._arrow_add_rows(foo=df).
-        elif len(kwargs) == 1:
-            name, data = kwargs.popitem()
-        # Raise error otherwise.
-        else:
-            raise StreamlitAPIException(
-                "Wrong number of arguments to add_rows()."
-                "Command requires exactly one dataset"
-            )
-
-        # When doing _arrow_add_rows on an element that does not already have data
-        # (for example, st.line_chart() without any args), call the original
-        # st.foo() element with new data instead of doing a _arrow_add_rows().
-        if (
-            "add_rows_metadata" in self._cursor.props
-            and self._cursor.props["add_rows_metadata"]
-            and self._cursor.props["add_rows_metadata"].last_index is None
-        ):
-            st_method = getattr(
-                self, self._cursor.props["add_rows_metadata"].chart_command
-            )
-            st_method(data, **kwargs)
-            return None
-
-        new_data, self._cursor.props["add_rows_metadata"] = _prep_data_for_add_rows(
-            data,
-            self._cursor.props["add_rows_metadata"],
-        )
-
-        msg = ForwardMsg_pb2.ForwardMsg()
-        msg.metadata.delta_path[:] = self._cursor.delta_path
-
-        import streamlit.elements.arrow as arrow_proto
-
-        default_uuid = str(hash(self._get_delta_path_str()))
-        arrow_proto.marshall(msg.delta.arrow_add_rows.data, new_data, default_uuid)
-
-        if name:
-            msg.delta.arrow_add_rows.name = name
-            msg.delta.arrow_add_rows.has_name = True
-
-        _enqueue_message(msg)
-
-        return self
-
 
 main_dg = DeltaGenerator(root_container=RootContainer.MAIN)
 sidebar_dg = DeltaGenerator(root_container=RootContainer.SIDEBAR, parent=main_dg)
@@ -702,42 +582,7 @@ def get_last_dg_added_to_context_stack() -> DeltaGenerator | None:
     return None
 
 
-def _prep_data_for_add_rows(
-    data: Data,
-    add_rows_metadata: AddRowsMetadata | None,
-) -> tuple[Data, AddRowsMetadata | None]:
-    if not add_rows_metadata:
-        if dataframe_util.is_pandas_styler(data):
-            # When calling add_rows on st.table or st.dataframe we want styles to
-            # pass through.
-            return data, None
-        return dataframe_util.convert_anything_to_pandas_df(data), None
-
-    # If add_rows_metadata is set, it indicates that the add_rows used called
-    # on a chart based on our built-in chart commands.
-
-    # For built-in chart commands we have to reshape the data structure
-    # otherwise the input data and the actual data used
-    # by vega_lite will be different, and it will throw an error.
-    from streamlit.elements.lib.built_in_chart_utils import prep_chart_data_for_add_rows
-
-    return prep_chart_data_for_add_rows(data, add_rows_metadata)
-
-
-def _enqueue_message(msg: ForwardMsg_pb2.ForwardMsg) -> None:
-    """Enqueues a ForwardMsg proto to send to the app."""
-    ctx = get_script_run_ctx()
-
-    if ctx is None:
-        raise NoSessionContext()
-
-    if ctx.current_fragment_id and msg.WhichOneof("type") == "delta":
-        msg.delta.fragment_id = ctx.current_fragment_id
-
-    ctx.enqueue(msg)
-
-
-def _writes_directly_to_sidebar(dg: DG) -> bool:
+def _writes_directly_to_sidebar(dg: DeltaGenerator) -> bool:
     in_sidebar = any(a._root_container == RootContainer.SIDEBAR for a in dg._ancestors)
     has_container = bool(len(list(dg._ancestor_block_types)))
     return in_sidebar and not has_container

streamlit/elements/arrow.py

@@ -18,7 +18,9 @@ import json
 from dataclasses import dataclass
 from typing import (
     TYPE_CHECKING,
+    Any,
     Final,
+    Hashable,
     Iterable,
     Literal,
     TypedDict,
@@ -43,14 +45,22 @@ from streamlit.elements.lib.policies import check_widget_policies
 from streamlit.elements.lib.utils import Key, to_key
 from streamlit.errors import StreamlitAPIException
 from streamlit.proto.Arrow_pb2 import Arrow as ArrowProto
+from streamlit.proto.ForwardMsg_pb2 import ForwardMsg
 from streamlit.runtime.metrics_util import gather_metrics
-from streamlit.runtime.scriptrunner import get_script_run_ctx
+from streamlit.runtime.scriptrunner.script_run_context import (
+    enqueue_message,
+    get_script_run_ctx,
+)
 from streamlit.runtime.state import WidgetCallback, register_widget
 from streamlit.runtime.state.common import compute_widget_id
 
 if TYPE_CHECKING:
+    from numpy import typing as npt
+    from pandas import DataFrame
+
     from streamlit.dataframe_util import Data
     from streamlit.delta_generator import DeltaGenerator
+    from streamlit.elements.lib.built_in_chart_utils import AddRowsMetadata
 
 
 SelectionMode: TypeAlias = Literal[
@@ -671,7 +681,7 @@ class ArrowMixin:
         >>> my_chart.add_rows(some_fancy_name=df2)  # <-- name used as keyword
 
         """
-        return self.dg._arrow_add_rows(data, **kwargs)
+        return _arrow_add_rows(self.dg, data, **kwargs)
 
     @property
     def dg(self) -> DeltaGenerator:
@@ -679,6 +689,140 @@ class ArrowMixin:
         return cast("DeltaGenerator", self)
 
 
+def _prep_data_for_add_rows(
+    data: Data,
+    add_rows_metadata: AddRowsMetadata | None,
+) -> tuple[Data, AddRowsMetadata | None]:
+    if not add_rows_metadata:
+        if dataframe_util.is_pandas_styler(data):
+            # When calling add_rows on st.table or st.dataframe we want styles to
+            # pass through.
+            return data, None
+        return dataframe_util.convert_anything_to_pandas_df(data), None
+
+    # If add_rows_metadata is set, it indicates that the add_rows used called
+    # on a chart based on our built-in chart commands.
+
+    # For built-in chart commands we have to reshape the data structure
+    # otherwise the input data and the actual data used
+    # by vega_lite will be different, and it will throw an error.
+    from streamlit.elements.lib.built_in_chart_utils import prep_chart_data_for_add_rows
+
+    return prep_chart_data_for_add_rows(data, add_rows_metadata)
+
+
+def _arrow_add_rows(
+    dg: DeltaGenerator,
+    data: Data = None,
+    **kwargs: (
+        DataFrame | npt.NDArray[Any] | Iterable[Any] | dict[Hashable, Any] | None
+    ),
+) -> DeltaGenerator | None:
+    """Concatenate a dataframe to the bottom of the current one.
+
+    Parameters
+    ----------
+    data : pandas.DataFrame, pandas.Styler, numpy.ndarray, Iterable, dict, or None
+        Table to concat. Optional.
+
+    **kwargs : pandas.DataFrame, numpy.ndarray, Iterable, dict, or None
+        The named dataset to concat. Optional. You can only pass in 1
+        dataset (including the one in the data parameter).
+
+    Example
+    -------
+    >>> import streamlit as st
+    >>> import pandas as pd
+    >>> import numpy as np
+    >>>
+    >>> df1 = pd.DataFrame(
+    ...     np.random.randn(50, 20), columns=("col %d" % i for i in range(20))
+    ... )
+    >>> my_table = st.table(df1)
+    >>>
+    >>> df2 = pd.DataFrame(
+    ...     np.random.randn(50, 20), columns=("col %d" % i for i in range(20))
+    ... )
+    >>> my_table.add_rows(df2)
+    >>> # Now the table shown in the Streamlit app contains the data for
+    >>> # df1 followed by the data for df2.
+
+    You can do the same thing with plots. For example, if you want to add
+    more data to a line chart:
+
+    >>> # Assuming df1 and df2 from the example above still exist...
+    >>> my_chart = st.line_chart(df1)
+    >>> my_chart.add_rows(df2)
+    >>> # Now the chart shown in the Streamlit app contains the data for
+    >>> # df1 followed by the data for df2.
+
+    And for plots whose datasets are named, you can pass the data with a
+    keyword argument where the key is the name:
+
+    >>> my_chart = st.vega_lite_chart(
+    ...     {
+    ...         "mark": "line",
+    ...         "encoding": {"x": "a", "y": "b"},
+    ...         "datasets": {
+    ...             "some_fancy_name": df1,  # <-- named dataset
+    ...         },
+    ...         "data": {"name": "some_fancy_name"},
+    ...     }
+    ... )
+    >>> my_chart.add_rows(some_fancy_name=df2)  # <-- name used as keyword
+
+    """
+    if dg._root_container is None or dg._cursor is None:
+        return dg
+
+    if not dg._cursor.is_locked:
+        raise StreamlitAPIException("Only existing elements can `add_rows`.")
+
+    # Accept syntax st._arrow_add_rows(df).
+    if data is not None and len(kwargs) == 0:
+        name = ""
+    # Accept syntax st._arrow_add_rows(foo=df).
+    elif len(kwargs) == 1:
+        name, data = kwargs.popitem()
+    # Raise error otherwise.
+    else:
+        raise StreamlitAPIException(
+            "Wrong number of arguments to add_rows()."
+            "Command requires exactly one dataset"
+        )
+
+    # When doing _arrow_add_rows on an element that does not already have data
+    # (for example, st.line_chart() without any args), call the original
+    # st.foo() element with new data instead of doing a _arrow_add_rows().
+    if (
+        "add_rows_metadata" in dg._cursor.props
+        and dg._cursor.props["add_rows_metadata"]
+        and dg._cursor.props["add_rows_metadata"].last_index is None
+    ):
+        st_method = getattr(dg, dg._cursor.props["add_rows_metadata"].chart_command)
+        st_method(data, **kwargs)
+        return None
+
+    new_data, dg._cursor.props["add_rows_metadata"] = _prep_data_for_add_rows(
+        data,
+        dg._cursor.props["add_rows_metadata"],
+    )
+
+    msg = ForwardMsg()
+    msg.metadata.delta_path[:] = dg._cursor.delta_path
+
+    default_uuid = str(hash(dg._get_delta_path_str()))
+    marshall(msg.delta.arrow_add_rows.data, new_data, default_uuid)
+
+    if name:
+        msg.delta.arrow_add_rows.name = name
+        msg.delta.arrow_add_rows.has_name = True
+
+    enqueue_message(msg)
+
+    return dg
+
+
 def marshall(proto: ArrowProto, data: Data, default_uuid: str | None = None) -> None:
     """Marshall pandas.DataFrame into an Arrow proto.
 

streamlit/elements/deck_gl_json_chart.py

@@ -16,7 +16,7 @@ from __future__ import annotations
 
 import hashlib
 import json
-from typing import TYPE_CHECKING, Any, Final, Mapping, cast
+from typing import TYPE_CHECKING, Any, Dict, Final, Mapping, cast
 
 from streamlit import config
 from streamlit.proto.DeckGlJsonChart_pb2 import DeckGlJsonChart as PydeckProto
@@ -158,7 +158,7 @@ def _get_pydeck_tooltip(pydeck_obj: Deck | None) -> dict[str, str] | None:
     # For details, see: https://github.com/visgl/deck.gl/pull/7125/files
     tooltip = getattr(pydeck_obj, "_tooltip", None)
     if tooltip is not None and isinstance(tooltip, dict):
-        return tooltip
+        return cast(Dict[str, str], tooltip)
 
     return None
 

streamlit/elements/image.py

@@ -34,6 +34,7 @@ from streamlit.errors import StreamlitAPIException
 from streamlit.proto.Image_pb2 import ImageList as ImageListProto
 from streamlit.runtime import caching
 from streamlit.runtime.metrics_util import gather_metrics
+from streamlit.type_util import NumpyShape
 
 if TYPE_CHECKING:
     from typing import Any
@@ -259,16 +260,17 @@ def _4d_to_list_3d(array: npt.NDArray[Any]) -> list[npt.NDArray[Any]]:
 
 
 def _verify_np_shape(array: npt.NDArray[Any]) -> npt.NDArray[Any]:
-    if len(array.shape) not in (2, 3):
+    shape: NumpyShape = array.shape
+    if len(shape) not in (2, 3):
         raise StreamlitAPIException("Numpy shape has to be of length 2 or 3.")
-    if len(array.shape) == 3 and array.shape[-1] not in (1, 3, 4):
+    if len(shape) == 3 and shape[-1] not in (1, 3, 4):
         raise StreamlitAPIException(
             "Channel can only be 1, 3, or 4 got %d. Shape is %s"
-            % (array.shape[-1], str(array.shape))
+            % (shape[-1], str(shape))
         )
 
     # If there's only one channel, convert is to x, y
-    if len(array.shape) == 3 and array.shape[-1] == 1:
+    if len(shape) == 3 and shape[-1] == 1:
         array = array[:, :, 0]
 
     return array
@@ -418,7 +420,7 @@ def image_to_url(
         )
 
         if channels == "BGR":
-            if len(image.shape) == 3:
+            if len(cast(NumpyShape, image.shape)) == 3:
                 image = image[:, :, [2, 1, 0]]
             else:
                 raise StreamlitAPIException(
@@ -509,7 +511,7 @@ def marshall_images(
     images: Sequence[AtomicImage]
     if isinstance(image, list):
         images = image
-    elif isinstance(image, np.ndarray) and len(image.shape) == 4:
+    elif isinstance(image, np.ndarray) and len(cast(NumpyShape, image.shape)) == 4:
         images = _4d_to_list_3d(image)
     else:
         images = [image]
@@ -520,7 +522,10 @@ def marshall_images(
         if isinstance(caption, str):
             captions = [caption]
         # You can pass in a 1-D Numpy array as captions.
-        elif isinstance(caption, np.ndarray) and len(caption.shape) == 1:
+        elif (
+            isinstance(caption, np.ndarray)
+            and len(cast(NumpyShape, caption.shape)) == 1
+        ):
             captions = caption.tolist()
         # If there are no captions then make the captions list the same size
         # as the images list.