Fix #464 add pksetdefault1 for setdefault convenience (#465)

* Fix #464 add pksetdefault1 for setdefault convenience

* review

pykern/pkcollections.py

@@ -211,35 +211,66 @@ def pknested_set(self, qualifiers, value):
         return self
 
     def pksetdefault(self, *args, **kwargs):
-        """Get value or set it, possibly after evaluating arg.
+        """Set values for keys if key not already in self
 
-        Must pass an even number of args or kwargs, but not both. Each pair
-        is interpreted as (key, value).
+        Accepts args or kwargs, but not both.
+        `args` must be an even number and are interpreted in (key, value) order.
+        `kwargs` are accepted as key=value.
 
-        If self does not have `key`, then it will be set. If `value` is a callable,
-        it will be called to get the value to set.
 
-        Values will be called if they are callable
+        If self does not have `key`, then it will be set to
+        `value`. If `key` is already in self, its value is not changed.
+         If `value` is a callable, it will be called, and
+        `key` will be set to the returned value.
 
         Args:
-            key (object): value to get or set
+            key (object): key that will be assigned `value` if not already set
             value (object): if callable, will be called, else verbatim
         Returns:
             object: self
         """
-        assert not (args and kwargs), "one of args or kwargs must be set, but not both"
+        if args and kwargs:
+            raise AssertionError("one of args or kwargs must be set, but not both")
         if args:
-            assert (
-                len(args) % 2 == 0
-            ), "args must be an even number (pairs of key, value)"
+            if len(args) % 2 != 0:
+                raise AssertionError(
+                    "args must be an even number (pairs of key, value)"
+                )
             i = zip(args[0::2], args[1::2])
         else:
             i = kwargs.items()
         for k, v in i:
-            if k not in self:
-                self[k] = v() if callable(v) else v
+            self.__pksetdefault_one(k, v)
         return self
 
+    def pksetdefault1(self, *args, **kwargs):
+        """Get value if exists else set
+
+        Accepts one key and one value, either as two positional args
+        (``key, value``); or one keyword arg (``key=value``).
+
+
+        If self does not have `key`, then it will be set to
+        `value`. If `value` is a callable, it will be called, and
+        `key` will be set to the returned value.
+
+        Args:
+            key (object): value to get or set
+            value (object): if callable, will be called, else verbatim
+        Returns:
+            object: ``self[key]``; either `value` if just set, or preexisting value
+
+        """
+        if args and kwargs:
+            raise AssertionError("one of args or kwargs must be set, but not both")
+        if args:
+            if len(args) == 2:
+                return self.__pksetdefault_one(*args)
+        elif len(kwargs) == 1:
+            for k, v in kwargs.items():
+                return self.__pksetdefault_one(k, v)
+        raise AssertionError("must pass exactly one key and value")
+
     def pkunchecked_nested_get(self, qualifiers, default=None):
         """Split key on dots or iterable and return nested get calls
 
@@ -265,6 +296,11 @@ def pkupdate(self, *args, **kwargs):
         super(PKDict, self).update(*args, **kwargs)
         return self
 
+    def __pksetdefault_one(self, key, value):
+        if key not in self:
+            self[key] = value() if callable(value) else value
+        return self[key]
+
 
 class PKDictNameError(NameError):
     """Raised when a key matches a builtin attribute in `dict`."""

tests/pkcollections_test.py

@@ -82,32 +82,6 @@ def test_dict():
         n.missing_attribute()
     with pkexcept(KeyError):
         n["missing key"]
-    pkeq(13, n.pksetdefault("d1", lambda: 13).d1)
-    pkeq(13, n.pksetdefault("d1", "already set").d1)
-    pkeq(n, n.pksetdefault("d1", "already set", "d2", 2, "d3", 3, "d4", 4))
-    pkeq(2, n.d2)
-    pkeq(3, n.d3)
-    pkeq(4, n.d4)
-    for i in "already set", 2, 3, 4:
-        with pkexcept(KeyError):
-            n[i]
-    with pkexcept(AssertionError):
-        n.pksetdefault("a", "b", "c")
-    n = PKDict()
-    pkeq(13, n.pksetdefault(d1=13).d1)
-    pkeq(13, n.pksetdefault(d1="already set").d1)
-    pkeq(n, n.pksetdefault(d1="already set", d2=2, d3=3, d4=4))
-    pkeq(2, n.d2)
-    pkeq(3, n.d3)
-    pkeq(4, n.d4)
-    for i in "already set", 2, 3, 4:
-        with pkexcept(KeyError):
-            n[i]
-    # ok to call empty kwargs
-    n.pksetdefault()
-    pkeq(4, n.pkdel("d4"))
-    pkeq(None, n.pkdel("d4"))
-    pkeq("x", n.pkdel("d4", "x"))
 
 
 def test_dict_copy():
@@ -124,6 +98,16 @@ def test_dict_copy():
     pkeq(id(n.b), id(n.b))
 
 
+def test_dict_pkdel():
+    from pykern.pkcollections import PKDict
+    from pykern.pkunit import pkok, pkeq, pkexcept
+
+    n = PKDict(d4=4)
+    pkeq(4, n.pkdel("d4"))
+    pkeq(None, n.pkdel("d4"))
+    pkeq("x", n.pkdel("d4", "x"))
+
+
 def test_dict_pknested_get():
     from pykern.pkcollections import PKDict
     from pykern.pkunit import pkeq, pkexcept
@@ -155,6 +139,55 @@ def test_dict_pknested_set():
     pkeq(3, n.pknested_get(["other"]))
 
 
+def test_dict_pksetdefault():
+    from pykern.pkcollections import PKDict
+    from pykern.pkunit import pkok, pkeq, pkexcept
+
+    n = PKDict()
+    pkeq(13, n.pksetdefault("d1", lambda: 13).d1)
+    pkeq(13, n.pksetdefault("d1", "already set").d1)
+    pkeq(n, n.pksetdefault("d1", "already set", "d2", 2, "d3", 3, "d4", 4))
+    pkeq(2, n.d2)
+    pkeq(3, n.d3)
+    pkeq(4, n.d4)
+    for i in "already set", 2, 3, 4:
+        with pkexcept(KeyError):
+            n[i]
+    with pkexcept(AssertionError):
+        n.pksetdefault("a", "b", "c")
+    n = PKDict()
+    pkeq(13, n.pksetdefault(d1=13).d1)
+    pkeq(13, n.pksetdefault(d1="already set").d1)
+    pkeq(n, n.pksetdefault(d1="already set", d2=2, d3=3, d4=4))
+    pkeq(2, n.d2)
+    pkeq(3, n.d3)
+    pkeq(4, n.d4)
+    for i in "already set", 2, 3, 4:
+        with pkexcept(KeyError):
+            n[i]
+    # ok to call empty kwargs
+    n.pksetdefault()
+
+
+def test_dict_pksetdefault1():
+    from pykern.pkcollections import PKDict
+    from pykern.pkunit import pkok, pkeq, pkexcept
+
+    n = PKDict()
+    pkeq(13, n.pksetdefault1("d1", lambda: 13))
+    pkeq(13, n.pksetdefault1("d1", "already set"))
+    pkeq("value", n.pksetdefault1(key1="value"))
+    pkeq("other", n.pksetdefault1(key2=lambda: "other"))
+    for a, k, r in (
+        (("d2", 2, "d3", 3), dict(), "too many args"),
+        ((), dict(d2=2, d3=3), "too many kwargs"),
+        ((), dict(), "no args"),
+        (("d3",), dict(), "one arg"),
+    ):
+        with pkexcept(AssertionError, r):
+            n.pksetdefault1(*a, **k)
+
+
 def test_dict_pkupdate():
     from pykern.pkcollections import PKDict
     from pykern.pkunit import pkeq