From 18a2c3c237014591d172284560546a2f0ac1a883 Mon Sep 17 00:00:00 2001 From: Dan Redding <125183946+dangotbanned@users.noreply.github.com> Date: Thu, 23 May 2024 20:17:20 +0100 Subject: [PATCH] feat: Adds `ChartType` type and type guard `is_chart_type`. Change PurePath to Path type hints (#3420) * feat: Type hints: Adds `ChartType` and guard `is_chart_type` I'm working on a library that uses `altair`, and looking to make it an optional dependency. During this process, I thought `ChartType` and `is_chart_type` might be a good fit in `altair` itself. I don't believe any existing `altair` code would directly benefit from these additions - but you could likely reduce some `isinstance` checks with [type narrowing](https://typing.readthedocs.io/en/latest/spec/narrowing.html) functions like `is_chart_type`. For example `altair/vegalite/v5/api.py` has **63** occurrences of `isinstance`. --- per [CONTRIBUTING.md](https://github.com/vega/altair/blob/main/CONTRIBUTING.md) > In particular, we welcome companion efforts from other visualization libraries to render the Vega-Lite specifications output by Altair. We see this portion of the effort as much bigger than Altair itself Towards this aim, it could be beneficial to provide `typing` constructs/tools for downstream use. Regarding my specfic use case, the below [Protocol](https://typing.readthedocs.io/en/latest/spec/protocol.html#runtime-checkable-decorator-and-narrowing-types-by-isinstance) - I believe - describes the concept of a `_/Chart`, for the purposes of writing to file: ```py from typing import Any, Protocol, runtime_checkable @runtime_checkable class AltairChart(Protocol): data: Any def copy(self, *args: Any, **kwargs: Any) -> Any: ... def save(self, *args: Any, **kwargs: Any) -> None: ... def to_dict(self, *args: Any, **kwargs: Any) -> dict[Any, Any]: ... def to_html(self, *args: Any, **kwargs: Any) -> str: ... ``` Due to the number of symbols in `altair` and heavy use of mixins, coming to this conclusion can be a stumbling block. However, exporting things like `ChartType`, `is_chart_type`, `AltairChart` could be a way to expose these shared behaviours - without requiring any changes to the core API. * fix: Type hints: Add missing `pathlib.Path` in `save` method annotation `TopLevelMixin.save` calls a function accepting `pathlib` objects, but the two didn't share the same `Union` for `fp`. When using `pyright`, the following warning is presented: > Argument of type "Path" cannot be assigned to parameter "fp" of type "str | IO[Unknown]" in function "save" > Type "Path" is incompatible with type "str | IO[Unknown]" > "Path" is incompatible with "str" > "Path" is incompatible with "IO[Unknown]" This fix avoids the need to use type ignores on the public facing API, as a user. * refactor: Simplify `save.py` path handling By normalizing `fp` in `save`, it simplifies the signatures of two other functions and removes the need for `str` specific suffix handling. * refactor: Replace duplicated `save` get encoding with a single expression Although there are branches it isn't needed, I think this is easier to read and maintain. Very low priority change, just something I noticed while reading. * docs: Add docstring for `is_chart_type` * fix: Removed unused `type: ignore` comment * build: Mark `TypeIs` as not a relevant attribute * build: Make new additions to `api` private Believe these should be public, but it seems like the issue at build-time is that I'd be extending the public API. * Bump version of typing_extensions so it includes TypeIs * Make ChartType and is_chart_type public. Use typing.get_args to reduce duplication of list of charts * Update API docs * refactor: Adds ruff `SIM101` and apply fixes As [requested](https://github.com/vega/altair/pull/3420/files/a4de6d1feb6f3f17c7f6d4c7738eb008ff33d17e#r1606780541). Seems to only occur in a few places, but all were autofix-able. * revert: Change `set` back to `list` for `format` As [requested](https://github.com/vega/altair/pull/3420/commits/99b4f81d456b2b05c97919e751152dcec1cbf00b#r1606785466) * style: Re-run ruff format * revert: Add `str` support back for "private" `save` functions An earlier commit assumed a minor optimization would be possible, by normalizing a `Union` early and reuse the narrowed type in "private" functions. @binste [comment](https://github.com/vega/altair/pull/3420/files#r1610440552) --------- Co-authored-by: Stefan Binder --- altair/__init__.py | 2 ++ altair/utils/_transformed_data.py | 4 ++-- altair/utils/save.py | 26 +++++++++----------------- altair/utils/schemapi.py | 2 +- altair/vegalite/v5/api.py | 23 ++++++++++++++++++++++- doc/user_guide/api.rst | 1 + pyproject.toml | 5 ++++- tools/schemapi/schemapi.py | 2 +- tools/update_init_file.py | 5 +++++ 9 files changed, 47 insertions(+), 23 deletions(-) diff --git a/altair/__init__.py b/altair/__init__.py index 073a1a53e..55feb3cdb 100644 --- a/altair/__init__.py +++ b/altair/__init__.py @@ -56,6 +56,7 @@ "Categorical", "Chart", "ChartDataType", + "ChartType", "Color", "ColorDatum", "ColorDef", @@ -580,6 +581,7 @@ "expr", "graticule", "hconcat", + "is_chart_type", "jupyter", "layer", "limit_rows", diff --git a/altair/utils/_transformed_data.py b/altair/utils/_transformed_data.py index 1d886eb5e..7a616b9c5 100644 --- a/altair/utils/_transformed_data.py +++ b/altair/utils/_transformed_data.py @@ -180,8 +180,8 @@ def name_views( List of the names of the charts and subcharts """ exclude = set(exclude) if exclude is not None else set() - if isinstance(chart, _chart_class_mapping[Chart]) or isinstance( - chart, _chart_class_mapping[FacetChart] + if isinstance( + chart, (_chart_class_mapping[Chart], _chart_class_mapping[FacetChart]) ): if chart.name not in exclude: if chart.name in (None, Undefined): diff --git a/altair/utils/save.py b/altair/utils/save.py index 11db77bf1..609486bc6 100644 --- a/altair/utils/save.py +++ b/altair/utils/save.py @@ -10,14 +10,14 @@ def write_file_or_filename( - fp: Union[str, pathlib.PurePath, IO], + fp: Union[str, pathlib.Path, 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): + if isinstance(fp, (str, pathlib.Path)): with open(file=fp, mode=mode, encoding=encoding) as f: f.write(content) else: @@ -25,14 +25,12 @@ def write_file_or_filename( def set_inspect_format_argument( - format: Optional[str], fp: Union[str, pathlib.PurePath, IO], inline: bool + format: Optional[str], fp: Union[str, pathlib.Path, IO], inline: bool ) -> str: """Inspect the format argument in the save function""" if format is None: - if isinstance(fp, str): - format = fp.split(".")[-1] - elif isinstance(fp, pathlib.PurePath): - format = fp.suffix.lstrip(".") + if isinstance(fp, (str, pathlib.Path)): + format = pathlib.Path(fp).suffix.lstrip(".") else: raise ValueError( "must specify file format: " @@ -71,7 +69,7 @@ def set_inspect_mode_argument( def save( chart, - fp: Union[str, pathlib.PurePath, IO], + fp: Union[str, pathlib.Path, IO], vega_version: Optional[str], vegaembed_version: Optional[str], format: Optional[Literal["json", "html", "png", "svg", "pdf"]] = None, @@ -140,7 +138,7 @@ def save( if json_kwds is None: json_kwds = {} - + encoding = kwargs.get("encoding", "utf-8") format = set_inspect_format_argument(format, fp, inline) # type: ignore[assignment] def perform_save(): @@ -152,9 +150,7 @@ def perform_save(): if format == "json": json_spec = json.dumps(spec, **json_kwds) - write_file_or_filename( - fp, json_spec, mode="w", encoding=kwargs.get("encoding", "utf-8") - ) + write_file_or_filename(fp, json_spec, mode="w", encoding=encoding) elif format == "html": if inline: kwargs["template"] = "inline" @@ -170,10 +166,7 @@ def perform_save(): **kwargs, ) write_file_or_filename( - fp, - mimebundle["text/html"], - mode="w", - encoding=kwargs.get("encoding", "utf-8"), + fp, mimebundle["text/html"], mode="w", encoding=encoding ) elif format in ["png", "svg", "pdf", "vega"]: mimebundle = spec_to_mimebundle( @@ -193,7 +186,6 @@ def perform_save(): elif format == "pdf": write_file_or_filename(fp, mimebundle["application/pdf"], mode="wb") else: - encoding = kwargs.get("encoding", "utf-8") write_file_or_filename( fp, mimebundle["image/svg+xml"], mode="w", encoding=encoding ) diff --git a/altair/utils/schemapi.py b/altair/utils/schemapi.py index 9325b3a70..10b2597e6 100644 --- a/altair/utils/schemapi.py +++ b/altair/utils/schemapi.py @@ -1193,7 +1193,7 @@ def _freeze(val): return frozenset((k, _freeze(v)) for k, v in val.items()) elif isinstance(val, set): return frozenset(map(_freeze, val)) - elif isinstance(val, list) or isinstance(val, tuple): + elif isinstance(val, (list, tuple)): return tuple(map(_freeze, val)) else: return val diff --git a/altair/vegalite/v5/api.py b/altair/vegalite/v5/api.py index 0bd9ae389..37a13c2ae 100644 --- a/altair/vegalite/v5/api.py +++ b/altair/vegalite/v5/api.py @@ -8,6 +8,8 @@ from toolz.curried import pipe as _pipe import itertools import sys +import pathlib +import typing from typing import cast, List, Optional, Any, Iterable, Union, Literal, IO # Have to rename it here as else it overlaps with schema.core.Type and schema.core.Dict @@ -30,6 +32,10 @@ from ...utils.data import DataType from ...utils.deprecation import AltairDeprecationWarning +if sys.version_info >= (3, 13): + from typing import TypeIs +else: + from typing_extensions import TypeIs if sys.version_info >= (3, 11): from typing import Self else: @@ -1138,7 +1144,7 @@ def open_editor(self, *, fullscreen: bool = False) -> None: def save( self, - fp: Union[str, IO], + fp: Union[str, pathlib.Path, IO], format: Optional[Literal["json", "html", "png", "svg", "pdf"]] = None, override_data_transformer: bool = True, scale_factor: float = 1.0, @@ -4098,3 +4104,18 @@ def graticule(**kwds): def sphere() -> core.SphereGenerator: """Sphere generator.""" return core.SphereGenerator(sphere=True) + + +ChartType = Union[ + Chart, RepeatChart, ConcatChart, HConcatChart, VConcatChart, FacetChart, LayerChart +] + + +def is_chart_type(obj: Any) -> TypeIs[ChartType]: + """Return `True` if the object is an Altair chart. This can be a basic chart + but also a repeat, concat, or facet chart. + """ + return isinstance( + obj, + typing.get_args(ChartType), + ) diff --git a/doc/user_guide/api.rst b/doc/user_guide/api.rst index 99aa5f69b..08013023f 100644 --- a/doc/user_guide/api.rst +++ b/doc/user_guide/api.rst @@ -152,6 +152,7 @@ API Functions condition graticule hconcat + is_chart_type layer param repeat diff --git a/pyproject.toml b/pyproject.toml index 36c95b317..a7ad687f7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -15,7 +15,7 @@ build-backend = "hatchling.build" name = "altair" authors = [{ name = "Vega-Altair Contributors" }] dependencies = [ - "typing_extensions>=4.0.1; python_version<\"3.11\"", + "typing_extensions>=4.10.0; python_version<\"3.13\"", "jinja2", # If you update the minimum required jsonschema version, also update it in build.yml "jsonschema>=3.0", @@ -164,6 +164,7 @@ exclude = [ ] [tool.ruff.lint] +extend-safe-fixes=["SIM101"] select = [ # flake8-bugbear "B", @@ -177,6 +178,8 @@ select = [ "F", # flake8-tidy-imports "TID", + # flake8-simplify + "SIM101" ] ignore = [ # Whitespace before ':' diff --git a/tools/schemapi/schemapi.py b/tools/schemapi/schemapi.py index 29da970ec..ce1268c04 100644 --- a/tools/schemapi/schemapi.py +++ b/tools/schemapi/schemapi.py @@ -1191,7 +1191,7 @@ def _freeze(val): return frozenset((k, _freeze(v)) for k, v in val.items()) elif isinstance(val, set): return frozenset(map(_freeze, val)) - elif isinstance(val, list) or isinstance(val, tuple): + elif isinstance(val, (list, tuple)): return tuple(map(_freeze, val)) else: return val diff --git a/tools/update_init_file.py b/tools/update_init_file.py index 4f342f8ef..e4fa65a86 100644 --- a/tools/update_init_file.py +++ b/tools/update_init_file.py @@ -21,6 +21,10 @@ cast, ) +if sys.version_info >= (3, 13): + from typing import TypeIs +else: + from typing_extensions import TypeIs if sys.version_info >= (3, 11): from typing import Self else: @@ -97,6 +101,7 @@ def _is_relevant_attribute(attr_name: str) -> bool: or attr is Protocol or attr is Sequence or attr is IO + or attr is TypeIs or attr_name == "TypingDict" or attr_name == "TypingGenerator" or attr_name == "ValueOrDatum"