From 6911314cd436c4bdeb100f0ccd7a1289246c873e Mon Sep 17 00:00:00 2001
From: Matthias Diener <mdiener@illinois.edu>
Date: Wed, 14 Feb 2024 09:58:52 -0600
Subject: [PATCH] support pyrsistent, immutables, constantdict in KeyBuilder

---
 .github/workflows/ci.yml             |  2 +-
 pytools/persistent_dict.py           |  3 +++
 pytools/test/test_persistent_dict.py | 38 +++++++++++++---------------
 3 files changed, 21 insertions(+), 22 deletions(-)

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 4149996..c29a01f 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -83,7 +83,7 @@ jobs:
                 # AK, 2020-12-13
                 rm pytools/mpiwrap.py
 
-                EXTRA_INSTALL="numpy frozendict immutabledict orderedsets"
+                EXTRA_INSTALL="numpy frozendict immutabledict orderedsets constantdict immutables pyrsistent"
                 curl -L -O https://gitlab.tiker.net/inducer/ci-support/raw/main/build-and-test-py-project.sh
                 . ./build-and-test-py-project.sh
 
diff --git a/pytools/persistent_dict.py b/pytools/persistent_dict.py
index b979214..421e02a 100644
--- a/pytools/persistent_dict.py
+++ b/pytools/persistent_dict.py
@@ -402,6 +402,9 @@ class KeyBuilder:
             (self.rec(self.new_hash(), (k, v)).digest() for k, v in key.items()))
 
     update_for_immutabledict = update_for_frozendict
+    update_for_constantdict = update_for_frozendict
+    update_for_PMap = update_for_frozendict  # noqa: N815
+    update_for_Map = update_for_frozendict  # noqa: N815
 
     # }}}
 
diff --git a/pytools/test/test_persistent_dict.py b/pytools/test/test_persistent_dict.py
index dd7f9a2..eb22ec6 100644
--- a/pytools/test/test_persistent_dict.py
+++ b/pytools/test/test_persistent_dict.py
@@ -421,33 +421,29 @@ def test_scalar_hashing():
     assert keyb(np.clongdouble(1.1+2.2j)) == keyb(np.clongdouble(1.1+2.2j))
 
 
-def test_frozendict_hashing():
-    pytest.importorskip("frozendict")
-    from frozendict import frozendict
-
-    keyb = KeyBuilder()
-
-    d = {"a": 1, "b": 2}
-
-    assert keyb(frozendict(d)) == keyb(frozendict(d))
-    assert keyb(frozendict(d)) != keyb(frozendict({"a": 1, "b": 3}))
-    assert keyb(frozendict(d)) == keyb(frozendict({"b": 2, "a": 1}))
-
-    with pytest.raises(TypeError):
-        keyb(d)
-
+@pytest.mark.parametrize("dict_impl", ("immutabledict", "frozendict",
+                                       "constantdict",
+                                       ("immutables", "Map"),
+                                       ("pyrsistent", "pmap")))
+def test_dict_hashing(dict_impl):
+    if isinstance(dict_impl, str):
+        dict_package = dict_impl
+        dict_class = dict_impl
+    else:
+        dict_package = dict_impl[0]
+        dict_class = dict_impl[1]
 
-def test_immutabledict_hashing():
-    pytest.importorskip("immutabledict")
-    from immutabledict import immutabledict
+    pytest.importorskip(dict_package)
+    import importlib
+    dc = getattr(importlib.import_module(dict_package), dict_class)
 
     keyb = KeyBuilder()
 
     d = {"a": 1, "b": 2}
 
-    assert keyb(immutabledict(d)) == keyb(immutabledict(d))
-    assert keyb(immutabledict(d)) != keyb(immutabledict({"a": 1, "b": 3}))
-    assert keyb(immutabledict(d)) == keyb(immutabledict({"b": 2, "a": 1}))
+    assert keyb(dc(d)) == keyb(dc(d))
+    assert keyb(dc(d)) != keyb(dc({"a": 1, "b": 3}))
+    assert keyb(dc(d)) == keyb(dc({"b": 2, "a": 1}))
 
 
 def test_frozenset_hashing():
-- 
GitLab