Various typing/dataclasses fixes

doc/conf.py

@@ -24,4 +24,10 @@
     "sympy": ("https://docs.sympy.org/dev/", None),
     "typing_extensions":
         ("https://typing-extensions.readthedocs.io/en/latest/", None),
+    "immutabledict":
+        ("https://immutabledict.corenting.fr/", None)
+}
+autodoc_type_aliases = {
+    "ExpressionT": "ExpressionT",
+    "ArithmeticExpressionT": "ArithmeticExpressionT",
 }

doc/primitives.rst

@@ -1,6 +1,7 @@
 Primitives (Basic Objects)
 ==========================
 
+.. automodule:: pymbolic.typing
 .. automodule:: pymbolic.primitives
 
 .. vim: sw=4

pymbolic/__init__.py

@@ -40,16 +40,20 @@
 
 from .polynomial import Polynomial
 
-from .primitives import Variable as var  # noqa: N813
-from .primitives import variables
-from .primitives import flattened_sum
-from .primitives import subscript
-from .primitives import flattened_product
-from .primitives import quotient
-from .primitives import linear_combination
-from .primitives import make_common_subexpression as cse
-from .primitives import make_sym_vector
-from .primitives import disable_subscript_by_getitem
+from .primitives import (Variable as var,  # noqa: N813
+    Variable,
+    Expression,
+    variables,
+    flattened_sum,
+    subscript,
+    flattened_product,
+    quotient,
+    linear_combination,
+    make_common_subexpression as cse,
+    make_sym_vector,
+    disable_subscript_by_getitem,
+    expr_dataclass,
+)
 from .parser import parse
 from .mapper.evaluator import evaluate
 from .mapper.evaluator import evaluate_kw
@@ -60,10 +64,18 @@
 from .mapper.distributor import distribute as expand
 from .mapper.distributor import distribute
 from .mapper.flattener import flatten
+from .typing import NumberT, ScalarT, ArithmeticExpressionT, ExpressionT, BoolT
 
 
 __all__ = (
+    "ArithmeticExpressionT",
+    "BoolT",
+    "Expression",
+    "ExpressionT",
+    "NumberT",
     "Polynomial",
+    "ScalarT",
+    "Variable",
     "compile",
     "compiler",
     "cse",
@@ -78,6 +90,7 @@
     "evaluate_kw",
     "evaluator",
     "expand",
+    "expr_dataclass",
     "flatten",
     "flattened_product",
     "flattened_sum",

pymbolic/mapper/__init__.py

@@ -49,7 +49,7 @@
     .. rubric:: Handling objects that don't declare mapper methods
 
     In particular, this includes many non-subclasses of
-    :class:`pymbolic.primitives.Expression`.
+    :class:`pymbolic.Expression`.
 
     .. automethod:: map_foreign
 
@@ -113,16 +113,16 @@ class UnsupportedExpressionError(ValueError):
 # {{{ mapper base
 
 class Mapper:
-    """A visitor for trees of :class:`pymbolic.primitives.Expression`
+    """A visitor for trees of :class:`pymbolic.Expression`
     subclasses. Each expression-derived object is dispatched to the
-    method named by the :attr:`pymbolic.primitives.Expression.mapper_method`
+    method named by the :attr:`pymbolic.Expression.mapper_method`
     attribute and if not found, the methods named by the class attribute
     *mapper_method* in the method resolution order of the object.
     """
 
     def handle_unsupported_expression(self, expr, *args, **kwargs):
         """Mapper method that is invoked for
-        :class:`pymbolic.primitives.Expression` subclasses for which a mapper
+        :class:`pymbolic.Expression` subclasses for which a mapper
         method does not exist in this mapper.
         """
 
@@ -496,9 +496,9 @@ def map_call_with_kwargs(self, expr, *args, **kwargs):
         parameters = tuple([
             self.rec(child, *args, **kwargs) for child in expr.parameters
             ])
-        kw_parameters = {
+        kw_parameters = immutabledict({
                 key: self.rec(val, *args, **kwargs)
-                for key, val in expr.kw_parameters.items()}
+                for key, val in expr.kw_parameters.items()})
 
         if (function is expr.function
             and all(child is orig_child for child, orig_child in

pymbolic/mapper/stringifier.py

@@ -84,11 +84,11 @@
 class StringifyMapper(Mapper):
     """A mapper to turn an expression tree into a string.
 
-    :class:`pymbolic.primitives.Expression.__str__` is often implemented using
+    :class:`pymbolic.Expression.__str__` is often implemented using
     this mapper.
 
-    When it encounters an unsupported :class:`pymbolic.primitives.Expression`
-    subclass, it calls its :meth:`pymbolic.primitives.Expression.make_stringifier`
+    When it encounters an unsupported :class:`pymbolic.Expression`
+    subclass, it calls its :meth:`pymbolic.Expression.make_stringifier`
     method to get a :class:`StringifyMapper` that potentially does.
     """