diff --git a/altair/utils/display.py b/altair/utils/display.py index f2bb99bad..ffe8f6ac7 100644 --- a/altair/utils/display.py +++ b/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 ) diff --git a/altair/utils/html.py b/altair/utils/html.py index c1084aeec..b2aeb5514 100644 --- a/altair/utils/html.py +++ b/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", diff --git a/altair/utils/mimebundle.py b/altair/utils/mimebundle.py index fa36ed4b8..4948904ed 100644 --- a/altair/utils/mimebundle.py +++ b/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'} diff --git a/altair/utils/save.py b/altair/utils/save.py index 797210b3f..97f86c710 100644 --- a/altair/utils/save.py +++ b/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}) diff --git a/altair/vegalite/v5/api.py b/altair/vegalite/v5/api.py index a5e280874..f0126ac38 100644 --- a/altair/vegalite/v5/api.py +++ b/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. diff --git a/altair/vegalite/v5/display.py b/altair/vegalite/v5/display.py index 714e17b42..93926e62b 100644 --- a/altair/vegalite/v5/display.py +++ b/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",