Add type annotations for plotly_chart

lib/streamlit/elements/plotly_chart.py

@@ -15,18 +15,38 @@
 """Streamlit support for Plotly charts."""
 
 import json
+import sys
 import urllib.parse
-from typing import cast
+from typing import Any, cast, Dict, List, Set, TYPE_CHECKING, Union
+
+if sys.version_info >= (3, 8):
+    from typing import Final, Literal
+else:
+    from typing_extensions import Final, Literal
+
+if sys.version_info >= (3, 10):
+    from typing import TypeAlias
+else:
+    from typing_extensions import TypeAlias
 
-import streamlit
 from streamlit.legacy_caching import caching
 from streamlit import type_util
 from streamlit.logger import get_logger
 from streamlit.proto.PlotlyChart_pb2 import PlotlyChart as PlotlyChartProto
 
-LOGGER = get_logger(__name__)
+if TYPE_CHECKING:
+    import matplotlib
+    import plotly.graph_objs as go
+    from plotly.basedatatypes import BaseFigure
+
+    from streamlit.delta_generator import DeltaGenerator
+
 
-SHARING_MODES = {
+LOGGER: Final = get_logger(__name__)
+
+SharingMode: TypeAlias = Literal["streamlit", "private", "public", "secret"]
+
+SHARING_MODES: Set[SharingMode] = {
     # This means the plot will be sent to the Streamlit app rather than to
     # Plotly.
     "streamlit",
@@ -37,15 +57,30 @@
     "secret",
 }
 
+_AtomicFigureOrData: TypeAlias = Union[
+    "go.Figure",
+    "go.Data",
+]
+FigureOrData: TypeAlias = Union[
+    _AtomicFigureOrData,
+    List[_AtomicFigureOrData],
+    # It is kind of hard to figure out exactly what kind of dict is supported
+    # here, as plotly hasn't embraced typing yet. This version is chosen to
+    # align with the docstring.
+    Dict[str, _AtomicFigureOrData],
+    "BaseFigure",
+    "matplotlib.figure.Figure",
+]
+
 
 class PlotlyMixin:
     def plotly_chart(
         self,
-        figure_or_data,
-        use_container_width=False,
-        sharing="streamlit",
-        **kwargs,
-    ):
+        figure_or_data: FigureOrData,
+        use_container_width: bool = False,
+        sharing: SharingMode = "streamlit",
+        **kwargs: Any,
+    ) -> "DeltaGenerator":
         """Display an interactive Plotly chart.
 
         Plotly is a charting library for Python. The arguments to this function
@@ -115,15 +150,24 @@ def plotly_chart(
         marshall(
             plotly_chart_proto, figure_or_data, use_container_width, sharing, **kwargs
         )
-        return self.dg._enqueue("plotly_chart", plotly_chart_proto)
+        return cast(
+            "DeltaGenerator",
+            self.dg._enqueue("plotly_chart", plotly_chart_proto),
+        )
 
     @property
-    def dg(self) -> "streamlit.delta_generator.DeltaGenerator":
+    def dg(self) -> "DeltaGenerator":
         """Get our DeltaGenerator."""
-        return cast("streamlit.delta_generator.DeltaGenerator", self)
+        return cast("DeltaGenerator", self)
 
 
-def marshall(proto, figure_or_data, use_container_width, sharing, **kwargs):
+def marshall(
+    proto: PlotlyChartProto,
+    figure_or_data: FigureOrData,
+    use_container_width: bool,
+    sharing: SharingMode,
+    **kwargs: Any,
+) -> None:
     """Marshall a proto with a Plotly spec.
 
     See DeltaGenerator.plotly_chart for docs.
@@ -166,10 +210,10 @@ def marshall(proto, figure_or_data, use_container_width, sharing, **kwargs):
 
 
 @caching.cache
-def _plot_to_url_or_load_cached_url(*args, **kwargs):
+def _plot_to_url_or_load_cached_url(*args: Any, **kwargs: Any) -> "go.Figure":
     """Call plotly.plot wrapped in st.cache.
 
-    This is so we don't unecessarily upload data to Plotly's SASS if nothing
+    This is so we don't unnecessarily upload data to Plotly's SASS if nothing
     changed since the previous upload.
     """
     try:
@@ -181,7 +225,7 @@ def _plot_to_url_or_load_cached_url(*args, **kwargs):
     return ply.plot(*args, **kwargs)
 
 
-def _get_embed_url(url):
+def _get_embed_url(url: str) -> str:
     parsed_url = urllib.parse.urlparse(url)
 
     # Plotly's embed URL is the normal URL plus ".embed".