feat: Improve the syntax for conditions with multiple predicates

altair/__init__.py

@@ -588,7 +588,6 @@
     "load_ipython_extension",
     "load_schema",
     "mixins",
-    "overload",
     "param",
     "parse_shorthand",
     "pipe",

altair/utils/schemapi.py

@@ -893,7 +893,7 @@ def to_dict(
         *,
         ignore: Optional[List[str]] = None,
         context: Optional[Dict[str, Any]] = None,
-    ) -> dict:
+    ) -> Dict[str, Any]:
         """Return a dictionary representation of the object
 
         Parameters

altair/vegalite/v5/api.py

@@ -1,3 +1,4 @@
+# ruff: noqa: F811
 import warnings
 
 import hashlib
@@ -10,7 +11,18 @@
 import sys
 import pathlib
 import typing
-from typing import cast, List, Optional, Any, Iterable, Union, Literal, IO
+from typing import (
+    cast,
+    List,
+    Optional,
+    Any,
+    Iterable,
+    Union,
+    Literal,
+    IO,
+    overload,
+    Tuple,
+)
 
 # Have to rename it here as else it overlaps with schema.core.Type and schema.core.Dict
 from typing import Type as TypingType
@@ -19,7 +31,7 @@
 from .schema import core, channels, mixins, Undefined, UndefinedType, SCHEMA_URL
 
 from .data import data_transformers
-from ... import utils, expr
+from ... import utils, expr  # noqa: F401 [No longer needed here, but keeping to prevent breakage elsewhere]
 from ...expr import core as _expr_core
 from .display import renderers, VEGALITE_VERSION, VEGAEMBED_VERSION, VEGA_VERSION
 from .theme import themes
@@ -345,6 +357,118 @@ def check_fields_and_encodings(parameter: Parameter, field_name: str) -> bool:
     return False
 
 
+# -------------------------------------------------------------------------
+# Tools for working with conditions
+_TestPredicateType = Union[str, _expr_core.Expression, core.PredicateComposition]
+"""TODO
+
+Item [(2)](https://vega.github.io/vega-lite/docs/condition.html) Specifying a `test` predicate: ...
+"""
+
+_PredicateType = Union[Parameter, core.Expr, TypingDict[str, Any], _TestPredicateType]
+"""Permitted types for `predicate`."""
+
+_DictOrStr = Union[TypingDict[str, Any], str]
+_DictOrSchema = Union[core.SchemaBase, TypingDict[str, Any]]
+
+_StatementType = Union[core.SchemaBase, _DictOrStr]
+"""Permitted types for `if_true`/`if_false`.
+
+In python terms:
+```py
+if _PredicateType:
+    return _StatementType
+elif _PredicateType:
+    return _StatementType
+else:
+    return _StatementType
+```
+"""
+
+_ConditionType = TypingDict[str, Union[_TestPredicateType, UndefinedType, Any]]
+"""Intermediate type representing a converted `_PredicateType`.
+
+Prior to parsing any `_StatementType`.
+"""
+
+_SelectionType = Union[core.SchemaBase, TypingDict[str, Union[_ConditionType, Any]]]
+"""Returned type of `alt.condition`.
+"""
+
+
+def _is_test_predicate(obj: Any) -> TypeIs[_TestPredicateType]:
+    return isinstance(obj, (str, _expr_core.Expression, core.PredicateComposition))
+
+
+def _get_predicate_expr(p: Parameter) -> Union[Any, UndefinedType]:
+    return getattr(p.param, "expr", Undefined)
+
+
+def _predicate_to_condition(
+    predicate: _PredicateType, **kwargs: Any
+) -> Tuple[_ConditionType, TypingDict[str, Any]]:
+    # - Empty is only popped when `Parameter`
+    # - Otherwise, its added to `if_true`, `if_false` later
+    # - Needs to be preserved in either case, as there are 2 sites,
+    #   where its used for updating a `dict`
+    condition: _ConditionType
+    if isinstance(predicate, Parameter):
+        if (
+            predicate.param_type == "selection"
+            or _get_predicate_expr(predicate) is Undefined
+        ):
+            condition = {"param": predicate.name}
+            if "empty" in kwargs:
+                condition["empty"] = kwargs.pop("empty")
+            elif isinstance(predicate.empty, bool):
+                condition["empty"] = predicate.empty
+        else:
+            condition = {"test": _get_predicate_expr(predicate)}
+    elif _is_test_predicate(predicate):
+        condition = {"test": predicate}
+    elif isinstance(predicate, dict):
+        condition = predicate
+    else:
+        msg = "condition predicate of type {}" "".format(type(predicate))
+        raise NotImplementedError(msg)
+    return condition, kwargs
+
+
+def _condition_to_selection(
+    condition: _ConditionType,
+    if_true: _StatementType,
+    if_false: _StatementType,
+    **kwargs: Any,
+) -> _SelectionType:
+    selection: _SelectionType
+    if isinstance(if_true, core.SchemaBase):
+        # NOTE: Removed an outdated comment
+        if_true = if_true.to_dict()
+    elif isinstance(if_true, str):
+        if isinstance(if_false, str):
+            msg = (
+                "A field cannot be used for both the `if_true` and `if_false` "
+                "values of a condition. "
+                "One of them has to specify a `value` or `datum` definition."
+            )
+            raise ValueError(msg)
+        else:
+            if_true = utils.parse_shorthand(if_true)
+            if_true.update(kwargs)
+    condition.update(if_true)
+    if isinstance(if_false, core.SchemaBase):
+        # For the selection, the channel definitions all allow selections
+        # already. So use this SchemaBase wrapper if possible.
+        selection = if_false.copy()
+        selection.condition = condition
+    elif isinstance(if_false, str):
+        selection = {"condition": condition, "shorthand": if_false}
+        selection.update(kwargs)
+    else:
+        selection = dict(condition=condition, **if_false)
+    return selection
+
+
 # ------------------------------------------------------------------------
 # Top-Level Functions
 
@@ -790,17 +914,35 @@ def binding_range(**kwargs):
     return core.BindRange(input="range", **kwargs)
 
 
+@overload
+def condition(
+    predicate: _PredicateType,
+    if_true: _StatementType,
+    if_false: core.SchemaBase,
+    **kwargs,
+) -> core.SchemaBase: ...
+@overload
+def condition(
+    predicate: _PredicateType, if_true: str, if_false: str, **kwargs
+) -> typing.NoReturn: ...
+@overload
+def condition(
+    predicate: _PredicateType, if_true: _DictOrSchema, if_false: _DictOrStr, **kwargs
+) -> TypingDict[str, Union[_ConditionType, Any]]: ...
+@overload
+def condition(
+    predicate: _PredicateType,
+    if_true: _DictOrStr,
+    if_false: TypingDict[str, Any],
+    **kwargs,
+) -> TypingDict[str, Union[_ConditionType, Any]]: ...
 # TODO: update the docstring
 def condition(
-    predicate: Union[
-        Parameter, str, expr.Expression, core.Expr, core.PredicateComposition, dict
-    ],
-    # Types of these depends on where the condition is used so we probably
-    # can't be more specific here.
-    if_true: Any,
-    if_false: Any,
+    predicate: _PredicateType,
+    if_true: _StatementType,
+    if_false: _StatementType,
     **kwargs,
-) -> Union[dict, core.SchemaBase]:
+) -> _SelectionType:
     """A conditional attribute or encoding
 
     Parameters
@@ -820,62 +962,8 @@ def condition(
     spec: dict or VegaLiteSchema
         the spec that describes the condition
     """
-    test_predicates = (str, expr.Expression, core.PredicateComposition)
-
-    condition: TypingDict[
-        str,
-        Union[
-            bool, str, _expr_core.Expression, core.PredicateComposition, UndefinedType
-        ],
-    ]
-    if isinstance(predicate, Parameter):
-        if (
-            predicate.param_type == "selection"
-            or getattr(predicate.param, "expr", Undefined) is Undefined
-        ):
-            condition = {"param": predicate.name}
-            if "empty" in kwargs:
-                condition["empty"] = kwargs.pop("empty")
-            elif isinstance(predicate.empty, bool):
-                condition["empty"] = predicate.empty
-        else:
-            condition = {"test": getattr(predicate.param, "expr", Undefined)}
-    elif isinstance(predicate, test_predicates):
-        condition = {"test": predicate}
-    elif isinstance(predicate, dict):
-        condition = predicate
-    else:
-        raise NotImplementedError(
-            "condition predicate of type {}" "".format(type(predicate))
-        )
-
-    if isinstance(if_true, core.SchemaBase):
-        # convert to dict for now; the from_dict call below will wrap this
-        # dict in the appropriate schema
-        if_true = if_true.to_dict()
-    elif isinstance(if_true, str):
-        if isinstance(if_false, str):
-            raise ValueError(
-                "A field cannot be used for both the `if_true` and `if_false` values of a condition. One of them has to specify a `value` or `datum` definition."
-            )
-        else:
-            if_true = utils.parse_shorthand(if_true)
-            if_true.update(kwargs)
-    condition.update(if_true)
-
-    selection: Union[dict, core.SchemaBase]
-    if isinstance(if_false, core.SchemaBase):
-        # For the selection, the channel definitions all allow selections
-        # already. So use this SchemaBase wrapper if possible.
-        selection = if_false.copy()
-        selection.condition = condition
-    elif isinstance(if_false, str):
-        selection = {"condition": condition, "shorthand": if_false}
-        selection.update(kwargs)
-    else:
-        selection = dict(condition=condition, **if_false)
-
-    return selection
+    condition, kwargs = _predicate_to_condition(predicate, **kwargs)
+    return _condition_to_selection(condition, if_true, if_false, **kwargs)
 
 
 # --------------------------------------------------------------------

tests/utils/test_utils.py

@@ -1,7 +1,7 @@
 import io
 import json
 import warnings
-
+import sys
 import numpy as np
 import pandas as pd
 import pytest
@@ -121,6 +121,9 @@ def test_sanitize_dataframe_arrow_columns():
 
 
 @pytest.mark.skipif(pa is None, reason="pyarrow not installed")
+@pytest.mark.skipif(
+    sys.platform == "win32", reason="Timezone database is not installed on Windows"
+)
 def test_sanitize_pyarrow_table_columns():
     # create a dataframe with various types
     df = pd.DataFrame(

tools/generate_api_docs.py

@@ -103,7 +103,7 @@ def api_functions() -> List[str]:
     altair_api_functions = [
         obj_name
         for obj_name in iter_objects(alt.api, restrict_to_type=types.FunctionType)  # type: ignore[attr-defined]
-        if obj_name != "cast"
+        if obj_name not in {"cast", "overload"}
     ]
     return sorted(altair_api_functions)
 

tools/update_init_file.py

@@ -19,6 +19,8 @@
     TypeVar,
     Union,
     cast,
+    overload,
+    Tuple,
 )
 
 if sys.version_info >= (3, 13):
@@ -102,6 +104,8 @@ def _is_relevant_attribute(attr_name: str) -> bool:
         or attr is Sequence
         or attr is IO
         or attr is TypeIs
+        or attr is overload
+        or attr is Tuple
         or attr_name == "TypingDict"
         or attr_name == "TypingGenerator"
         or attr_name == "ValueOrDatum"