feat: Reimplement expr as a class that is understood by IDEs

altair/expr/consts.py

@@ -1,7 +1,5 @@
 from __future__ import annotations
 
-from .core import ConstExpression
-
 
 CONST_LISTING = {
     "NaN": "not a number (same as JavaScript literal NaN)",
@@ -14,16 +12,3 @@
     "SQRT2": "the square root of 2 (alias to Math.SQRT1_2)",
     "PI": "the transcendental number pi (alias to Math.PI)",
 }
-
-NAME_MAP: dict[str, str] = {}
-
-
-def _populate_namespace():
-    globals_ = globals()
-    for name, doc in CONST_LISTING.items():
-        py_name = NAME_MAP.get(name, name)
-        globals_[py_name] = ConstExpression(name, doc)
-        yield py_name
-
-
-__all__ = list(_populate_namespace())

altair/expr/core.py

@@ -210,9 +210,8 @@ def __repr__(self):
 
 
 class ConstExpression(Expression):
-    def __init__(self, name, doc) -> None:
-        self.__doc__ = f"""{name}: {doc}"""
-        super().__init__(name=name, doc=doc)
+    def __init__(self, name) -> None:
+        super().__init__(name=name)
 
     def __repr__(self) -> str:
         return str(self.name)

altair/expr/funcs.py

@@ -1,10 +1,4 @@
 from __future__ import annotations
-from typing import TYPE_CHECKING
-from .core import FunctionExpression
-
-if TYPE_CHECKING:
-    from typing import Iterator
-
 
 FUNCTION_LISTING = {
     "isArray": r"Returns true if _value_ is an array, false otherwise.",
@@ -171,27 +165,3 @@
 
 # This maps vega expression function names to the Python name
 NAME_MAP = {"if": "if_"}
-
-
-class ExprFunc:
-    def __init__(self, name, doc) -> None:
-        self.name = name
-        self.doc = doc
-        self.__doc__ = f"""{name}(*args)\n    {doc}"""
-
-    def __call__(self, *args) -> FunctionExpression:
-        return FunctionExpression(self.name, args)
-
-    def __repr__(self) -> str:
-        return f"<function expr.{self.name}(*args)>"
-
-
-def _populate_namespace() -> Iterator[str]:
-    globals_ = globals()
-    for name, doc in FUNCTION_LISTING.items():
-        py_name = NAME_MAP.get(name, name)
-        globals_[py_name] = ExprFunc(name, doc)
-        yield py_name
-
-
-__all__ = list(_populate_namespace())

altair/vegalite/v5/__init__.py

@@ -2,7 +2,7 @@
 from .schema import *
 from .api import *
 
-from ...expr import datum, expr
+from altair.expr.core import datum
 
 from .display import VegaLite, renderers
 from .compiler import vegalite_compilers

doc/user_guide/api.rst

@@ -637,6 +637,7 @@ API Utility Classes
    :toctree: generated/api-cls/
    :nosignatures:
 
+   expr
    When
    Then
    ChainedWhen

tests/expr/test_expr.py

@@ -1,12 +1,39 @@
+from __future__ import annotations
+
 import operator
+import sys
+from inspect import classify_class_attrs, getmembers
+from typing import Any, Iterator
 
 import pytest
-
-from altair import expr
-from altair import datum
-from altair import ExprRef
 from jsonschema.exceptions import ValidationError
 
+from altair import ExprRef, datum, expr
+from altair.expr import _ConstExpressionType
+
+# This maps vega expression function names to the Python name
+VEGA_REMAP = {"if_": "if"}
+
+
+def _is_property(obj: Any, /) -> bool:
+    return isinstance(obj, property)
+
+
+def _get_classmethod_names(tp: type[Any], /) -> Iterator[str]:
+    for m in classify_class_attrs(tp):
+        if m.kind == "class method" and m.defining_class is tp:
+            yield m.name
+
+
+def _remap_classmethod_names(tp: type[Any], /) -> Iterator[tuple[str, str]]:
+    for name in _get_classmethod_names(tp):
+        yield VEGA_REMAP.get(name, name), name
+
+
+def _get_property_names(tp: type[Any], /) -> Iterator[str]:
+    for nm, _ in getmembers(tp, _is_property):
+        yield nm
+
 
 def test_unary_operations():
     OP_MAP = {"-": operator.neg, "+": operator.pos}
@@ -59,22 +86,34 @@ def test_abs():
     assert repr(z) == "abs(datum.xxx)"
 
 
-def test_expr_funcs():
+@pytest.mark.parametrize(("veganame", "methodname"), _remap_classmethod_names(expr))
+def test_expr_funcs(veganame: str, methodname: str):
     """test all functions defined in expr.funcs"""
-    name_map = {val: key for key, val in expr.funcs.NAME_MAP.items()}
-    for funcname in expr.funcs.__all__:
-        func = getattr(expr, funcname)
-        z = func(datum.xxx)
-        assert repr(z) == f"{name_map.get(funcname, funcname)}(datum.xxx)"
+    func = getattr(expr, methodname)
+    z = func(datum.xxx)
+    assert repr(z) == f"{veganame}(datum.xxx)"
 
 
-def test_expr_consts():
+@pytest.mark.parametrize("constname", _get_property_names(_ConstExpressionType))
+def test_expr_consts(constname: str):
     """Test all constants defined in expr.consts"""
-    name_map = {val: key for key, val in expr.consts.NAME_MAP.items()}
-    for constname in expr.consts.__all__:
-        const = getattr(expr, constname)
-        z = const * datum.xxx
-        assert repr(z) == f"({name_map.get(constname, constname)} * datum.xxx)"
+
+    const = getattr(expr, constname)
+    z = const * datum.xxx
+    assert repr(z) == f"({constname} * datum.xxx)"
+
+
+@pytest.mark.parametrize("constname", _get_property_names(_ConstExpressionType))
+def test_expr_consts_immutable(constname: str):
+    """Ensure e.g `alt.expr.PI = 2` is prevented."""
+    if sys.version_info >= (3, 11):
+        pattern = f"property {constname!r}.+has no setter"
+    elif sys.version_info >= (3, 10):
+        pattern = f"can't set attribute {constname!r}"
+    else:
+        pattern = "can't set attribute"
+    with pytest.raises(AttributeError, match=pattern):
+        setattr(expr, constname, 2)
 
 
 def test_json_reprs():
@@ -127,7 +166,7 @@ def test_expression_function_nostring():
     # expr() can only work with str otherwise
     # should raise a SchemaValidationError
     with pytest.raises(ValidationError):
-        expr(2 * 2)
+        expr(2 * 2)  # pyright: ignore
 
     with pytest.raises(ValidationError):
-        expr(["foo", "bah"])
+        expr(["foo", "bah"])  # pyright: ignore

tests/vegalite/v5/test_api.py

@@ -593,13 +593,15 @@ def test_when_expressions_inside_parameters() -> None:
     )
     cond = when_then_otherwise["condition"][0]
     otherwise = when_then_otherwise["value"]
-    expected = alt.expr(alt.expr.if_(alt.datum.b >= 0, 10, -20))
-    actual = alt.expr(alt.expr.if_(cond["test"], cond["value"], otherwise))
+    expected = alt.expr.if_(alt.datum.b >= 0, 10, -20)
+    actual = alt.expr.if_(cond["test"], cond["value"], otherwise)
     assert expected == actual
 
-    text_conditioned = bar.mark_text(align="left", baseline="middle", dx=actual).encode(
-        text="b"
-    )
+    text_conditioned = bar.mark_text(
+        align="left",
+        baseline="middle",
+        dx=alt.expr(actual),  # type: ignore[arg-type]
+    ).encode(text="b")
 
     chart = bar + text_conditioned
     chart.to_dict()
@@ -719,12 +721,14 @@ def test_selection_to_dict():
 
 
 def test_selection_expression():
+    from altair.expr.core import Expression
+
     selection = alt.selection_point(fields=["value"])
 
     assert isinstance(selection.value, alt.SelectionExpression)
     assert selection.value.to_dict() == {"expr": f"{selection.name}.value"}
 
-    assert isinstance(selection["value"], alt.expr.Expression)
+    assert isinstance(selection["value"], Expression)
     assert selection["value"].to_dict() == f"{selection.name}['value']"
 
     magic_attr = "__magic__"

tests/vegalite/v5/test_params.py

@@ -117,13 +117,15 @@ def test_parameter_naming():
 
 
 def test_selection_expression():
+    from altair.expr.core import Expression
+
     data = pd.DataFrame([{"a": "A", "b": 28}])
 
     sel = alt.selection_point(fields=["b"])
     se = sel.b | 300
 
     assert isinstance(se, alt.SelectionExpression)
-    assert isinstance(se.expr, alt.expr.core.Expression)
+    assert isinstance(se.expr, Expression)
 
     c = (
         alt.Chart(data)

tools/generate_api_docs.py

@@ -115,9 +115,8 @@ def api_functions() -> list[str]:
 
 
 def api_classes() -> list[str]:
-    # classes defined in `api` and returned by `API Functions`,
-    # but not covered in other groups
-    return ["When", "Then", "ChainedWhen"]
+    # Part of the Public API, but are not inherited from `vega-lite`.
+    return ["expr", "When", "Then", "ChainedWhen"]
 
 
 def lowlevel_wrappers() -> list[str]: