Repository: radiasoft/pykern

Issue:
add PKDict.pknested_set

Discussion:

Patch:
diff --git a/pykern/pkcollections.py b/pykern/pkcollections.py
index e195a63..e2683e8 100644
--- a/pykern/pkcollections.py
+++ b/pykern/pkcollections.py
@@ -168,7 +168,7 @@ class PKDict(dict):
         Throws KeyError if the dictionary key doesn't exist.
 
         Args:
-            qualifiers (str or iterable):
+            qualifiers (str or iterable): see above
 
         Returns:
             object: value of element
@@ -186,6 +186,31 @@ class PKDict(dict):
                 raise
         return d
 
+    def pknested_set(self, qualifiers, value):
+        """Set nested location identified by `qualifiers` to `value`
+
+        If `qualifiers` is a str, will split on dots. Otherwise, will be iterated.
+
+        If an element does not exist, creates it as a PKDict.
+
+        DOES NOT SUPPORT ints (lists) at this time.
+
+        Args:
+            qualifiers (str or iterable): see above
+            value (any): assigned to qualifier's location in self
+
+        Returns:
+            object: self
+        """
+        q = qualifiers.split(".") if isinstance(qualifiers, str) else list(qualifiers)
+        d = self
+        for k in q[:-1]:
+            if k not in d:
+                d[k] = PKDict()
+            d = d[k]
+        d[q[-1]] = value
+        return self
+
     def pksetdefault(self, *args, **kwargs):
         """Get value or set it, possibly after evaluating arg.
 


Test patch:
diff --git a/tests/pkcollections_test.py b/tests/pkcollections_test.py
index 1d800e6..4778460 100644
--- a/tests/pkcollections_test.py
+++ b/tests/pkcollections_test.py
@@ -144,6 +144,17 @@ def test_dict_pknested_get():
     pkeq("done", n.pknested_get("one.1.last"))
 
 
+def test_dict_pknested_set():
+    from pykern.pkcollections import PKDict
+    from pykern.pkunit import pkeq, pkexcept
+
+    n = PKDict(one=PKDict())
+    n.pknested_set("one.two", 2)
+    n.pknested_set("other", 3)
+    pkeq(2, n.pknested_get("one.two"))
+    pkeq(3, n.pknested_get(["other"]))
+
+
 def test_dict_pkupdate():
     from pykern.pkcollections import PKDict
     from pykern.pkunit import pkeq