Merge pull request #500 from GavinHuttley/develop

ENH: UnionDict setitem works like conventional dict too
cogent3 · Jan 20, 2020 · fee26cd · fee26cd
2 parents e4d9427 + bb60ad5
commit fee26cd
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 0 deletions.
diff --git a/src/cogent3/util/union_dict.py b/src/cogent3/util/union_dict.py
@@ -54,6 +54,12 @@ def __setattr__(self, key, value):
 
         self.update({key: value})
 
+    def __setitem__(self, key, value):
+        if isinstance(value, dict):
+            value = UnionDict(value)
+
+        self.update({key: value})
+
     def __or__(self, other):
         self.union(other)
         return self

diff --git a/tests/test_util/test_union_dict.py b/tests/test_util/test_union_dict.py
@@ -99,6 +99,18 @@ def test_get_subattr(self):
         self.assertEqual(d._getsubattr_([], "d"), UnionDict({"e": 5, "f": 6}))
         self.assertEqual(d._getsubattr_(["d"], "e"), 5)
 
+    def test_setitem(self):
+        """should work via property or key"""
+        d = UnionDict()
+        d.a = 23
+        d.b = dict(c=42)
+        self.assertEqual(d.a, 23)
+        self.assertEqual(d["a"], 23)
+        self.assertEqual(d.b, dict(c=42))
+        self.assertEqual(d.b.c, 42)
+        self.assertEqual(d["b"], dict(c=42))
+        self.assertIsInstance(d.b, UnionDict)
+
 
 if __name__ == "__main__":
     main()