Merge pull request #1408 from compas-dev/fix-serialisation-colordict

Fix serialisation of color dicts

Author: tomvanmele

CHANGELOG.md

@@ -9,6 +9,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+* Added key conversion map to `compas.colors.ColorDict` to avoid serialisation problems with tuple keys when used in combination with edges.
+
 ### Changed
 
 * Fixed bug in `VolMesh.delete_cell`.

src/compas/colors/colordict.py

@@ -24,16 +24,23 @@ class ColorDict(Data):
 
     """
 
+    KEYMAP = {
+        int: lambda x: str(x),
+        tuple: lambda x: ",".join(map(str, sorted(x))),
+        list: lambda x: ",".join(map(str, sorted(x))),
+    }
+
     @property
     def __data__(self):
-        return {"default": self.default, "dict": self._dict}
+        return {
+            "default": self.default,
+            "dict": self._dict,
+        }
 
     @classmethod
     def __from_data__(cls, data):
         colordict = cls(data["default"])
-        # note: this is specific to color dicts for vertices of meshes
-        # perhaps the color dict needs to be subclassed per scene object type
-        colordict.update({int(key): value for key, value in data["dict"].items()})
+        colordict.update(data["dict"])
         return colordict
 
     def __init__(self, default, name=None):
@@ -54,14 +61,19 @@ def default(self, default):
             default = Color.coerce(default)
         self._default = default
 
+    def keymapper(self, key):
+        if key.__class__ in self.KEYMAP:
+            return self.KEYMAP[key.__class__](key)
+        return key
+
     def __getitem__(self, key):
-        return self._dict.get(key, self.default)
+        return self._dict.get(self.keymapper(key), self.default)
 
     def __setitem__(self, key, value):
-        self._dict[key] = Color.coerce(value)
+        self._dict[self.keymapper(key)] = Color.coerce(value)
 
     def __delitem__(self, key):
-        del self._dict[key]
+        del self._dict[self.keymapper(key)]
 
     def __iter__(self):
         return iter(self._dict)
@@ -70,7 +82,7 @@ def __len__(self):
         return len(self._dict)
 
     def __contains__(self, key):
-        return key in self._dict
+        return self.keymapper(key) in self._dict
 
     def items(self):
         return self._dict.items()
@@ -82,7 +94,7 @@ def values(self):
         return self._dict.values()
 
     def get(self, key, default=None):
-        return self._dict.get(key, default or self.default)
+        return self._dict.get(self.keymapper(key), default or self.default)
 
     def clear(self):
         """Clear the previously stored items.
@@ -108,4 +120,4 @@ def update(self, other):
 
         """
         for key, value in other.items():
-            self[key] = value
+            self[self.keymapper(key)] = value

tests/compas/colors/test_colordict.py

@@ -0,0 +1,47 @@
+import pytest
+import json
+import compas
+
+from compas.colors import Color
+from compas.colors import ColorDict
+
+
+def test_colordict():
+    cd = ColorDict(Color.red())
+    assert cd.default == Color.red()
+
+    cd = ColorDict((255, 0, 0))
+    assert cd.default == Color.red()
+
+    cd = ColorDict((1.0, 0.0, 0.0))
+    assert cd.default == Color.red()
+
+
+def test_colordict_keys():
+    cd = ColorDict(Color.red())
+    cd[1] = Color.blue()
+    cd[(1, 0)] = Color.green()
+
+    assert cd[1] == Color.blue()
+    assert cd["1"] == Color.blue()
+
+    assert cd[(1, 0)] == Color.green()
+    assert cd["0,1"] == Color.green()
+
+    assert cd["1,0"] == Color.red()
+
+
+def test_colordict_json():
+    cd1 = ColorDict(Color.red())
+    cd1[1] = Color.blue()
+    cd1[(1, 0)] = Color.green()
+
+    cd2 = compas.json_loads(compas.json_dumps(cd1))
+
+    assert cd2[1] == Color.blue()
+    assert cd2["1"] == Color.blue()
+
+    assert cd2[(1, 0)] == Color.green()
+    assert cd2["0,1"] == Color.green()
+
+    assert cd2["1,0"] == Color.red()