diff --git a/pykern/pkcollections.py b/pykern/pkcollections.py index 062ed5c7..acdfabc8 100644 --- a/pykern/pkcollections.py +++ b/pykern/pkcollections.py @@ -117,25 +117,20 @@ def pksetdefault(self, *args, **kwargs): key (object): value to get or set value (object): if callable, will be called, else verbatim Returns: - object: self if multiple pairs else returns evaluated value + object: self """ assert bool(args) != bool(kwargs), \ 'one of args or kwargs must be set, but not both' - r = None if args: assert len(args) % 2 == 0, \ 'args must be an even number (pairs of key, value)' i = zip(args[0::2], args[1::2]) - if len(args) <= 2: - r = args[0] else: i = kwargs.items() - if len(kwargs) <= 1: - r = list(kwargs.keys())[0] for k, v in i: if k not in self: self[k] = v() if callable(v) else v - return self[r] if r else self + return self def pkupdate(self, *args, **kwargs): """Call `dict.update` and return ``self``. diff --git a/tests/pkcollections_test.py b/tests/pkcollections_test.py index c7dfaacf..124d42fd 100644 --- a/tests/pkcollections_test.py +++ b/tests/pkcollections_test.py @@ -69,8 +69,8 @@ def test_dict(): n.missing_attribute() with pkexcept(KeyError): n['missing key'] - pkeq(13, n.pksetdefault('d1', lambda: 13)) - pkeq(13, n.pksetdefault('d1', 'already set')) + 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) @@ -81,8 +81,8 @@ def test_dict(): with pkexcept(AssertionError): n.pksetdefault('a', 'b', 'c') n = PKDict() - pkeq(13, n.pksetdefault(d1=13)) - pkeq(13, n.pksetdefault(d1='already set')) + 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)