Type hint utils/save.py and utils/mimebundle.py (#3248)

* Type hint mimebundle.py

* Type hint save.py

* Fix type hint for py38

altair/utils/display.py

@@ -210,6 +210,8 @@ def output_div(self) -> str:
     def __call__(self, spec: dict, **metadata) -> Dict[str, str]:
         kwargs = self.kwargs.copy()
         kwargs.update(metadata)
-        return spec_to_mimebundle(
+        # To get proper return value type, would need to write complex
+        # overload signatures for spec_to_mimebundle based on `format`
+        return spec_to_mimebundle(  # type: ignore[return-value]
             spec, format="html", output_div=self.output_div, **kwargs
         )

altair/utils/html.py

@@ -213,8 +213,8 @@
 def spec_to_html(
     spec: dict,
     mode: str,
-    vega_version: str,
-    vegaembed_version: str,
+    vega_version: Optional[str],
+    vegaembed_version: Optional[str],
     vegalite_version: Optional[str] = None,
     base_url: str = "https://cdn.jsdelivr.net/npm",
     output_div: str = "vis",

altair/utils/mimebundle.py

@@ -1,3 +1,5 @@
+from typing import Literal, Optional, Union, cast, Tuple
+
 from .deprecation import AltairDeprecationWarning
 from .html import spec_to_html
 from ._importers import import_vl_convert
@@ -6,15 +8,15 @@
 
 
 def spec_to_mimebundle(
-    spec,
-    format,
-    mode=None,
-    vega_version=None,
-    vegaembed_version=None,
-    vegalite_version=None,
-    engine=None,
+    spec: dict,
+    format: Literal["html", "json", "png", "svg", "pdf", "vega", "vega-lite"],
+    mode: Optional[Literal["vega-lite"]] = None,
+    vega_version: Optional[str] = None,
+    vegaembed_version: Optional[str] = None,
+    vegalite_version: Optional[str] = None,
+    engine: Optional[Literal["vl-convert", "altair_saver"]] = None,
     **kwargs,
-):
+) -> Union[dict, Tuple[dict, dict]]:
     """Convert a vega-lite specification to a mimebundle
 
     The mimebundle type is controlled by the ``format`` argument, which can be
@@ -54,18 +56,20 @@ def spec_to_mimebundle(
     if mode != "vega-lite":
         raise ValueError("mode must be 'vega-lite'")
 
+    internal_mode: Literal["vega-lite", "vega"] = mode
     if using_vegafusion():
         spec = compile_with_vegafusion(spec)
-        mode = "vega"
+        internal_mode = "vega"
 
     if format in ["png", "svg", "pdf", "vega"]:
+        format = cast(Literal["png", "svg", "pdf", "vega"], format)
         return _spec_to_mimebundle_with_engine(
-            spec, format, mode, engine=engine, **kwargs
+            spec, format, internal_mode, engine=engine, **kwargs
         )
     if format == "html":
         html = spec_to_html(
             spec,
-            mode=mode,
+            mode=internal_mode,
             vega_version=vega_version,
             vegaembed_version=vegaembed_version,
             vegalite_version=vegalite_version,
@@ -84,7 +88,12 @@ def spec_to_mimebundle(
     )
 
 
-def _spec_to_mimebundle_with_engine(spec, format, mode, **kwargs):
+def _spec_to_mimebundle_with_engine(
+    spec: dict,
+    format: Literal["png", "svg", "pdf", "vega"],
+    mode: Literal["vega-lite", "vega"],
+    **kwargs,
+) -> Union[dict, Tuple[dict, dict]]:
     """Helper for Vega-Lite to mimebundle conversions that require an engine
 
     Parameters
@@ -183,7 +192,10 @@ def _spec_to_mimebundle_with_engine(spec, format, mode, **kwargs):
         )
 
 
-def _validate_normalize_engine(engine, format):
+def _validate_normalize_engine(
+    engine: Optional[Literal["vl-convert", "altair_saver"]],
+    format: Literal["png", "svg", "pdf", "vega"],
+) -> str:
     """Helper to validate and normalize the user-provided engine
 
     engine : {None, 'vl-convert', 'altair_saver'}

altair/utils/save.py

@@ -1,13 +1,19 @@
 import json
 import pathlib
 import warnings
+from typing import IO, Union, Optional, Literal
 
 from .mimebundle import spec_to_mimebundle
 from ..vegalite.v5.data import data_transformers
 from altair.utils._vegafusion_data import using_vegafusion
 
 
-def write_file_or_filename(fp, content, mode="w", encoding=None):
+def write_file_or_filename(
+    fp: Union[str, pathlib.PurePath, IO],
+    content: Union[str, bytes],
+    mode: str = "w",
+    encoding: Optional[str] = None,
+) -> None:
     """Write content to fp, whether fp is a string, a pathlib Path or a
     file-like object"""
     if isinstance(fp, str) or isinstance(fp, pathlib.PurePath):
@@ -17,7 +23,9 @@ def write_file_or_filename(fp, content, mode="w", encoding=None):
         fp.write(content)
 
 
-def set_inspect_format_argument(format, fp, inline):
+def set_inspect_format_argument(
+    format: Optional[str], fp: Union[str, pathlib.PurePath, IO], inline: bool
+) -> str:
     """Inspect the format argument in the save function"""
     if format is None:
         if isinstance(fp, str):
@@ -36,7 +44,12 @@ def set_inspect_format_argument(format, fp, inline):
     return format
 
 
-def set_inspect_mode_argument(mode, embed_options, spec, vegalite_version):
+def set_inspect_mode_argument(
+    mode: Optional[Literal["vega-lite"]],
+    embed_options: dict,
+    spec: dict,
+    vegalite_version: Optional[str],
+) -> Literal["vega-lite"]:
     """Inspect the mode argument in the save function"""
     if mode is None:
         if "mode" in embed_options:
@@ -57,20 +70,20 @@ def set_inspect_mode_argument(mode, embed_options, spec, vegalite_version):
 
 def save(
     chart,
-    fp,
-    vega_version,
-    vegaembed_version,
-    format=None,
-    mode=None,
-    vegalite_version=None,
-    embed_options=None,
-    json_kwds=None,
-    webdriver=None,
-    scale_factor=1,
-    engine=None,
-    inline=False,
+    fp: Union[str, pathlib.PurePath, IO],
+    vega_version: Optional[str],
+    vegaembed_version: Optional[str],
+    format: Optional[Literal["json", "html", "png", "svg", "pdf"]] = None,
+    mode: Optional[Literal["vega-lite"]] = None,
+    vegalite_version: Optional[str] = None,
+    embed_options: Optional[dict] = None,
+    json_kwds: Optional[dict] = None,
+    webdriver: Optional[Literal["chrome", "firefox"]] = None,
+    scale_factor: float = 1,
+    engine: Optional[Literal["vl-convert", "altair_saver"]] = None,
+    inline: bool = False,
     **kwargs,
-):
+) -> None:
     """Save a chart to file in a variety of formats
 
     Supported formats are [json, html, png, svg, pdf]
@@ -121,7 +134,7 @@ def save(
     if embed_options is None:
         embed_options = {}
 
-    format = set_inspect_format_argument(format, fp, inline)
+    format = set_inspect_format_argument(format, fp, inline)  # type: ignore[assignment]
 
     def perform_save():
         spec = chart.to_dict(context={"pre_transform": False})

altair/vegalite/v5/api.py

@@ -1128,10 +1128,10 @@ def save(
         # normally trigger a MaxRowsError
         if override_data_transformer:
             with data_transformers.disable_max_rows():
-                result = save(**kwds)
+                save(**kwds)
         else:
-            result = save(**kwds)
-        return result
+            save(**kwds)
+        return
 
     # Fallback for when rendering fails; the full repr is too long to be
     # useful in nearly all cases.

altair/vegalite/v5/display.py

@@ -59,7 +59,9 @@ def json_renderer(spec: dict, **metadata) -> DefaultRendererReturnType:
 
 
 def png_renderer(spec: dict, **metadata) -> Dict[str, bytes]:
-    return spec_to_mimebundle(
+    # To get proper return value type, would need to write complex
+    # overload signatures for spec_to_mimebundle based on `format`
+    return spec_to_mimebundle(  # type: ignore[return-value]
         spec,
         format="png",
         mode="vega-lite",
@@ -71,7 +73,9 @@ def png_renderer(spec: dict, **metadata) -> Dict[str, bytes]:
 
 
 def svg_renderer(spec: dict, **metadata) -> Dict[str, str]:
-    return spec_to_mimebundle(
+    # To get proper return value type, would need to write complex
+    # overload signatures for spec_to_mimebundle based on `format`
+    return spec_to_mimebundle(  # type: ignore[return-value]
         spec,
         format="svg",
         mode="vega-lite",