DOC: show ClassVar and default argument use

ComPWA · Dec 18, 2023 · 53f8a64 · 53f8a64
1 parent 0434d22
commit 53f8a64
Show file tree

Hide file tree

Showing 2 changed files with 79 additions and 0 deletions.
diff --git a/docs/usage/sympy.ipynb b/docs/usage/sympy.ipynb
@@ -162,6 +162,56 @@
     "Math(aslatex({e: e.evaluate() for e in [rho_expr, q_expr]}))"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Class variables and default arguments to instance arguments are also supported:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from __future__ import annotations\n",
+    "\n",
+    "from typing import Any, ClassVar\n",
+    "\n",
+    "\n",
+    "@unevaluated_expression\n",
+    "class FunkyPower(sp.Expr):\n",
+    "    x: Any\n",
+    "    m: int = 1\n",
+    "    default_return: ClassVar[sp.Expr | None] = None\n",
+    "    _latex_repr_ = R\"f_{{{m}}}\\left({x}\\right)\"\n",
+    "\n",
+    "    def evaluate(self) -> sp.Expr | None:\n",
+    "        if self.default_return is None:\n",
+    "            return self.x**self.m\n",
+    "        return self.default_return\n",
+    "\n",
+    "\n",
+    "x = sp.Symbol(\"x\")\n",
+    "exprs = (\n",
+    "    FunkyPower(x),\n",
+    "    FunkyPower(x, 2),\n",
+    "    FunkyPower(x, m=3),\n",
+    ")\n",
+    "Math(aslatex({e: e.doit() for e in exprs}))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "FunkyPower.default_return = sp.Rational(0.5)\n",
+    "Math(aslatex({e: e.doit() for e in exprs}))"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},

diff --git a/tests/sympy/test_decorator.py b/tests/sympy/test_decorator.py
@@ -134,6 +134,35 @@ def evaluate(self) -> sp.Expr:
     assert y_expr.doit() == 5**3
 
 
+def test_unevaluated_expression_classvar_symbol():
+    @unevaluated_expression
+    class FunkyPower(sp.Expr):
+        x: Any
+        m: int = 1
+        default_return: ClassVar[float | None] = None
+
+        def evaluate(self) -> sp.Expr:
+            if self.default_return is None:
+                return self.x**self.m
+            return sp.sympify(self.default_return)
+
+    x = sp.Symbol("x")
+    exprs = (
+        FunkyPower(x),
+        FunkyPower(x, 2),
+        FunkyPower(x, m=3),
+    )
+    assert exprs[0].doit() == x
+    assert exprs[1].doit() == x**2
+    assert exprs[2].doit() == x**3
+
+    half = sp.Rational(1, 2)
+    FunkyPower.default_return = half
+    assert exprs[0].doit() == half
+    assert exprs[1].doit() == half
+    assert exprs[2].doit() == half
+
+
 def test_unevaluated_expression_callable():
     @unevaluated_expression(implement_doit=False)
     class Squared(sp.Expr):