Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: Improve the syntax for conditions with multiple predicates #3427

Merged
merged 122 commits into from
Jul 18, 2024
Merged
Show file tree
Hide file tree
Changes from 13 commits
Commits
Show all changes
122 commits
Select commit Hold shift + click to select a range
303f784
fix: add overloads for `condition`
dangotbanned May 29, 2024
866a50e
build: regen `__all__` excluding `@overload`
dangotbanned May 29, 2024
7ddd214
fix: add statement overloads for `condition`
dangotbanned May 29, 2024
2b3c2f8
refactor: add `typing.Tuple` to `api` and `update_init_file`
dangotbanned May 30, 2024
61ab923
feat(typing): amend and extend aliases for conditions
dangotbanned May 30, 2024
a33efa3
feat(typing): update `condition` signature, overloads to use new aliases
dangotbanned May 30, 2024
79af12d
refactor: extract condition utility functions
dangotbanned May 30, 2024
d901679
refactor: extract condition parsing functions
dangotbanned May 30, 2024
03a372c
refactor: add `ruff` ignores to keep `expr` in `api`
dangotbanned May 30, 2024
8d455d1
refactor: replaced contents of `condition` with extracted functions
dangotbanned May 30, 2024
8df67e1
style: rerun `ruff`
dangotbanned May 30, 2024
aa35acf
test: Skip tests on Win that require a tz database
dangotbanned May 30, 2024
747ca6b
fix(typing): Correct `SchemaBase.to_dict` return type
dangotbanned May 31, 2024
14f6ea1
perf: Build base `units` for `utils.parse_shorthand` once
dangotbanned May 31, 2024
f20949d
feat: align both sides of `if_true`/ `if_false` parsing
dangotbanned May 31, 2024
9d202e5
test: Update single test that relied on intermediate shorthand
dangotbanned May 31, 2024
278f4f4
feat(typing): add `_is_statement_type` guard
dangotbanned May 31, 2024
113a793
docs: add reference to `predicate`
dangotbanned May 31, 2024
d6fea15
feat(typing): adds `_Conditions` intermediate representation
dangotbanned May 31, 2024
9c89bce
feat: Adds initial `when-then-otherwise` implementation
dangotbanned May 31, 2024
a09c13c
test: Add tests for `when-then-otherwise`
dangotbanned May 31, 2024
e7be069
fix(typing): Use `typing.List`
dangotbanned May 31, 2024
1ad06f4
fix(typing): fix(typing): Correct SchemaBase.to_dict return type
dangotbanned May 31, 2024
61279cd
fix(typing): Use safer `typing_extensions.runtime_checkable` unless `…
dangotbanned May 31, 2024
cf92402
perf: Don't recheck type after positive case in `infer_encoding_types`
dangotbanned May 31, 2024
22e89a8
feat: adds `_Then.to_dict`
dangotbanned May 31, 2024
dfc14de
feat: allow `_Then` as an encoding input
dangotbanned May 31, 2024
328338a
test: Add test for `_Then` as encoding condition
dangotbanned May 31, 2024
430898f
fix(typing): Remove `UndefinedType` from `_ConditionType`
dangotbanned Jun 2, 2024
5f49715
refactor: Simplify `_predicate_to_condition`
dangotbanned Jun 2, 2024
460444c
refactor: removed outdated comments
dangotbanned Jun 2, 2024
c8469ed
fix(typing): Correct `alt.value` return annotation
dangotbanned Jun 2, 2024
4b8353b
refactor: moved comment
dangotbanned Jun 2, 2024
f1d9dc2
feat(typing): extends aliases and guards
dangotbanned Jun 2, 2024
38d425f
feat: extend and refactor `when-then-otherwise` functionality
dangotbanned Jun 2, 2024
80c32d3
fix(typing): exclude `typing_extensions.Protocol` and short `TypeVar`…
dangotbanned Jun 2, 2024
7b5f711
test: update and extend tests for `when-then-otherwise`
dangotbanned Jun 2, 2024
b962c48
fix: add item to `pyproject.toml` to silence import errors for `pylan…
dangotbanned Jun 2, 2024
3a1c640
test: check output type on all `when-then-otherwise` type errors
dangotbanned Jun 2, 2024
7b671db
test: update tests, examples using deprecated `pandas` parameters
dangotbanned Jun 2, 2024
5f631de
feat: Support additional predicate types in `when`, `condition`
dangotbanned Jun 3, 2024
fbd9205
feat: adds conversion functions between `alt.condition`/`when-then-ot…
dangotbanned Jun 3, 2024
9af55c2
feat: Adds config to `alt.value` wrapping in `when-then-otherwise`
dangotbanned Jun 3, 2024
80fbcbd
test: replace now-removed `wrap_value` parameter
dangotbanned Jun 3, 2024
5da5622
test: additional `when-then-otherwise` tests
dangotbanned Jun 3, 2024
c994bf8
fix(typing): satisfy mypy
dangotbanned Jun 3, 2024
8c0f45a
revert: 7b671db85d344fd6e0ca2d2a25a862d84118988c
dangotbanned Jun 3, 2024
989b1f5
feat(typing): adds `_FieldEqualType` for when constraints
dangotbanned Jun 16, 2024
7d41219
docs, feat: Add draft doc for `alt.when` and make public
dangotbanned Jun 16, 2024
80a0812
docs: Replace markdown syntax with rst directives
dangotbanned Jun 16, 2024
bfa6d51
Merge branch 'vega:main' into condition-multiple
dangotbanned Jun 19, 2024
237c874
sync
dangotbanned Jun 19, 2024
f4b51b8
Merge branch 'condition-multiple' of https://github.com/dangotbanned/…
dangotbanned Jun 19, 2024
0a070c9
refactor: Remove unused `expr` import in `api`
dangotbanned Jun 19, 2024
b85223f
feat(typing): Preserve wrapper type in `alt.condition`
dangotbanned Jun 19, 2024
30c1ad1
fix(typing): remove now-unused ignores
dangotbanned Jun 19, 2024
95093d7
fix: Simplify complex `expr` import behaviour
dangotbanned Jun 19, 2024
55cc301
fix: silence `unused-ignore` for `test_api` until `Chart.encode` fix …
dangotbanned Jun 19, 2024
62ad8d1
docs: refine `alt.when` doc
dangotbanned Jun 22, 2024
9cf84f9
docs: Improve `_When` docs
dangotbanned Jun 22, 2024
58e0054
feat: make `alt.When` public
dangotbanned Jun 22, 2024
fe8ab25
test: Update test references to use public `alt.When`
dangotbanned Jun 22, 2024
e159ea0
docs: Fix some user guide sphinx errors
dangotbanned Jun 22, 2024
8e92953
Merge remote-tracking branch 'origin/main' into condition-multiple
dangotbanned Jun 28, 2024
910f493
revert(typing): re-enable mypy unused-ignore
dangotbanned Jun 28, 2024
2276d50
test: better utilise `pytest.raises(match=...)`
dangotbanned Jun 28, 2024
591a55b
refactor: UX-first improvements to imports, annotations
dangotbanned Jun 28, 2024
0be41e8
docs: Add doc for `_Then.otherwise`
dangotbanned Jun 29, 2024
11841d0
docs: Add doc for `_Then.when`
dangotbanned Jun 29, 2024
b48a812
feat: make `Then`, `ChainedWhen` public
dangotbanned Jun 29, 2024
8c1e36b
docs: Add doc for `ChainedWhen.then`, misc fixes
dangotbanned Jun 29, 2024
5baaa34
docs: add `api-cls` tree
dangotbanned Jun 29, 2024
801cccc
refactor: remove `seq_as_lit` option
dangotbanned Jun 30, 2024
25ac2fa
Merge branch 'main' into condition-multiple
dangotbanned Jun 30, 2024
b238aef
Merge branch 'main' into condition-multiple
dangotbanned Jul 2, 2024
5f7d949
chore: remove type ignore comments
dangotbanned Jul 2, 2024
f89990f
revert: remove condition -> expr conversion experiment
dangotbanned Jul 3, 2024
c217b4f
fix: `mypy` error
dangotbanned Jul 3, 2024
8f66b7a
Merge branch 'main' into condition-multiple
dangotbanned Jul 3, 2024
d12a741
feat: un-special case `Then`
dangotbanned Jul 6, 2024
87b6014
fix: exclude `_typing` imports from `__all__`
dangotbanned Jul 6, 2024
6f06b80
refactor(typing): replace `dict` w/ `Map` when possible
dangotbanned Jul 6, 2024
1292962
feat: redesign `when-then-otherwise` to account for condition restric…
dangotbanned Jul 6, 2024
b110f65
test: update test to use `Then`'s getattr
dangotbanned Jul 6, 2024
f33a9cd
test: adds `test_when_stress`, `test_when_condition_parity`
dangotbanned Jul 6, 2024
f95be1a
ci: fix `PT001` rule inversion conflict
dangotbanned Jul 6, 2024
31bad3f
feat(DRAFT): add experimental `_str_as`
dangotbanned Jul 6, 2024
d0784f8
refactor: `str_as_lit` -> `str_as`, invert default
dangotbanned Jul 7, 2024
fbd9088
Merge branch 'main' into condition-multiple
dangotbanned Jul 7, 2024
c8f7d6c
refactor: update `str_as` error
dangotbanned Jul 8, 2024
8f5d352
refactor: rename `kwargs` -> `kwds`
dangotbanned Jul 8, 2024
49c4d24
refactor(typing): use `_FieldEqualType` more consistently
dangotbanned Jul 8, 2024
38d8a8e
refactor(typing): consistently annotate `TypeAlias`s
dangotbanned Jul 8, 2024
bd26c07
refactor(typing): remove lesser used aliases
dangotbanned Jul 8, 2024
6245a8b
refactor(typing): replace lingering `UndefinedType`
dangotbanned Jul 8, 2024
5e64fcb
refactor: inline `_str_as`, add note for `PEP728`
dangotbanned Jul 8, 2024
a3531fb
test: remove `alt.value` from example-based tests
dangotbanned Jul 8, 2024
03d7ca5
Merge branch 'main' into condition-multiple
dangotbanned Jul 9, 2024
f8d7529
Merge branch 'main' into condition-multiple
dangotbanned Jul 13, 2024
c0afbad
Merge branch 'main' into condition-multiple
dangotbanned Jul 13, 2024
86880c9
merge remote
dangotbanned Jul 15, 2024
e712b61
Merge branch 'condition-multiple' of https://github.com/dangotbanned/…
dangotbanned Jul 15, 2024
d607c70
Merge branch 'main' into condition-multiple
dangotbanned Jul 15, 2024
f992324
Merge branch 'main' into condition-multiple
dangotbanned Jul 16, 2024
0f862ad
Merge branch 'main' into condition-multiple
dangotbanned Jul 18, 2024
290da86
refactor: remove `str_as` and `alt.value(statement)` wrapping
dangotbanned Jul 18, 2024
3fa4a3f
refactor(typing): remove/reorganize aliases and guards
dangotbanned Jul 18, 2024
528ca16
docs: minor consistency update to `_ComposablePredicateType`
dangotbanned Jul 18, 2024
159bcf1
chore: exclude `TypeAliasType` from `__all__`
dangotbanned Jul 18, 2024
500f7e1
test: update existing `when` tests
dangotbanned Jul 18, 2024
db473c3
test: adds tests @mattijn authored during review
dangotbanned Jul 18, 2024
de08676
docs: fix lowercase
dangotbanned Jul 18, 2024
633736d
docs: remove `str_as` references
dangotbanned Jul 18, 2024
0f54057
revert: remove unrelated `expr` changes
dangotbanned Jul 18, 2024
0ed9eab
docs(typing): add example to `_OneOrSeq`
dangotbanned Jul 18, 2024
fc6ce16
revert: Remove unrelated `pyproject.toml` changes
dangotbanned Jul 18, 2024
4bc4732
chore: remove TODO comment
dangotbanned Jul 18, 2024
f0bc7b8
refactor: Define `SHORTHAND_KEYS` in `utils.core`
dangotbanned Jul 18, 2024
33e17fa
docs: Update `_TestPredicateType` link
dangotbanned Jul 18, 2024
59af293
refactor: remove single use `_is_composable_type` guard
dangotbanned Jul 18, 2024
c70d535
fix: More useful error message for `alt.[condition|when](predicate)`
dangotbanned Jul 18, 2024
87cd9bd
docs: another consistency update
dangotbanned Jul 18, 2024
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 0 additions & 1 deletion altair/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -588,7 +588,6 @@
"load_ipython_extension",
"load_schema",
"mixins",
"overload",
"param",
"parse_shorthand",
"pipe",
Expand Down
2 changes: 1 addition & 1 deletion altair/utils/schemapi.py
Original file line number Diff line number Diff line change
Expand Up @@ -893,7 +893,7 @@ def to_dict(
*,
ignore: Optional[List[str]] = None,
context: Optional[Dict[str, Any]] = None,
) -> dict:
) -> Dict[str, Any]:
dangotbanned marked this conversation as resolved.
Show resolved Hide resolved
"""Return a dictionary representation of the object

Parameters
Expand Down
220 changes: 154 additions & 66 deletions altair/vegalite/v5/api.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
# ruff: noqa: F811
import warnings

import hashlib
Expand All @@ -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
Expand All @@ -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
dangotbanned marked this conversation as resolved.
Show resolved Hide resolved
from .display import renderers, VEGALITE_VERSION, VEGAEMBED_VERSION, VEGA_VERSION
from .theme import themes
Expand Down Expand Up @@ -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]
Copy link
Member Author

@dangotbanned dangotbanned Jun 2, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

core.Expr was in the union of alt.condition(predicate), but the internal logic would raise a NotImplementedError.
Keeping it around for now, but either the annotation or internals should be updated.

"""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

Expand Down Expand Up @@ -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
Expand All @@ -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)


# --------------------------------------------------------------------
Expand Down
5 changes: 4 additions & 1 deletion tests/utils/test_utils.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
import io
import json
import warnings

import sys
import numpy as np
import pandas as pd
import pytest
Expand Down Expand Up @@ -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(
Expand Down
2 changes: 1 addition & 1 deletion tools/generate_api_docs.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down
4 changes: 4 additions & 0 deletions tools/update_init_file.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,8 @@
TypeVar,
Union,
cast,
overload,
Tuple,
)

if sys.version_info >= (3, 13):
Expand Down Expand Up @@ -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"
Expand Down