allow defining a different EqualityMapper in Expression

Author: alexfikl

pymbolic/geometric_algebra/mapper.py

@@ -179,7 +179,7 @@ def map_nabla_component(self, expr):
         return {expr}
 
     def map_derivative_source(self, expr):
-        return {expr, self.rec(expr.operand)}
+        return {expr} | self.rec(expr.operand)
 
 
 class NablaComponentToUnitVector(EvaluationMapper):

pymbolic/mapper/__init__.py

@@ -104,8 +104,8 @@ def handle_unsupported_expression(self, expr, *args, **kwargs):
         """
 
         raise UnsupportedExpressionError(
-                "{} cannot handle expressions of type {}".format(
-                    type(self), type(expr)))
+                "'{}' cannot handle expressions of type '{}'".format(
+                    type(self).__name__, type(expr).__name__))
 
     def __call__(self, expr, *args, **kwargs):
         """Dispatch *expr* to its corresponding mapper method. Pass on

pymbolic/mapper/equality.py

@@ -22,38 +22,56 @@
 
 from typing import Any, Dict, Tuple
 
-from pymbolic.mapper import Mapper
+from pymbolic.mapper import Mapper, UnsupportedExpressionError
 from pymbolic.primitives import Expression
 
 
 class EqualityMapper(Mapper):
     def __init__(self) -> None:
         self._ids_to_result: Dict[Tuple[int, int], bool] = {}
 
-    def __call__(self, expr: Any, other: Any) -> bool:
+    def rec(self, expr: Any, other: Any) -> bool:
         key = (id(expr), id(other))
-        if key in self._ids_to_result:
-            return self._ids_to_result[key]
-
-        if expr is other:
-            result = True
-        elif expr.__class__ != other.__class__:
-            result = False
-        else:
-            try:
-                method = getattr(self, expr.mapper_method)
-            except AttributeError:
-                if isinstance(expr, Expression):
-                    return self.handle_unsupported_expression(expr, other)
-                else:
-                    return self.map_foreign(expr, other)
+
+        try:
+            result = self._ids_to_result[key]
+        except KeyError:
+            if expr is other:
+                result = True
+            elif expr.__class__ != other.__class__:
+                result = False
+            elif hash(expr) != hash(other):
+                result = False
             else:
-                result = method(expr, other)
+                try:
+                    method = getattr(self, expr.mapper_method)
+                except AttributeError:
+                    if isinstance(expr, Expression):
+                        result = self.handle_unsupported_expression(expr, other)
+                    else:
+                        result = self.map_foreign(expr, other)
+                else:
+                    result = method(expr, other)
+
+            self._ids_to_result[key] = result
 
-        self._ids_to_result[key] = result
         return result
 
-    rec = __call__
+    def __call__(self, expr: Any, other: Any) -> bool:
+        return self.rec(expr, other)
+
+    def handle_unsupported_expression(self, expr, other) -> bool:
+        eq = expr.make_equality_mapper()
+        if type(self) == type(eq):
+            raise UnsupportedExpressionError(
+                    "'{}' cannot handle expressions of type '{}'".format(
+                        type(self).__name__, type(expr).__name__))
+
+        eq._ids_to_result = self._ids_to_result
+        return eq(expr, other)
+
+    def map_constant(self, expr, other) -> bool:
+        return expr == other
 
     def map_variable(self, expr, other) -> bool:
         return expr.name == other.name

pymbolic/primitives.py

@@ -505,8 +505,7 @@ def __eq__(self, other):
         Subclasses should generally not override this method, but instead
         provide an implementation of :meth:`is_equal`.
         """
-        from pymbolic.mapper.equality import EqualityMapper
-        return EqualityMapper()(self, other)
+        return self.make_equality_mapper()(self, other)
 
     def __ne__(self, other):
         return not self.__eq__(other)
@@ -539,9 +538,18 @@ def __setstate__(self, state):
 
     # {{{ hash/equality backend
 
+    def make_equality_mapper(self):
+        from pymbolic.mapper.equality import EqualityMapper
+        return EqualityMapper()
+
     def is_equal(self, other):
-        return (type(other) == type(self)
-                and self.__getinitargs__() == other.__getinitargs__())
+        from warnings import warn
+        warn("'Expression.is_equal' is deprecated and will be removed in 2023. "
+                "To customize the equality check, subclass 'EqualityMapper' "
+                "and overwrite 'Expression.make_equality_mapper'",
+                DeprecationWarning, stacklevel=2)
+
+        return self.make_equality_mapper()(self, other)
 
     def get_hash(self):
         return hash((type(self).__name__,) + self.__getinitargs__())