Copying a Dictionary and Replace its Values

Question

I have a custom dictionary where the values are numpy arrays. I would like to have a convenience function that creates a copy of its content except that the values are changed. If no value is passed, then the default None is assigned to the keys, else a value is set. The function can also work with a callable function, and then the resulting value is obtained by passing the length of the numpy arrays.

Inspired by what I found here, this is what I have so far:

import sys
import numpy as np

class CustomDict(dict):

    def strip(self, value=None):

        if callable(value):
            tmp = dict.fromkeys(self.keys())
            for k,v in self.items():
                tmp[k] = value(len(v))
            return tmp

        return dict.fromkeys(self.keys(), value)


def main(argv=()):

    tmp1 = CustomDict()

    tmp1['n1'] = np.array([[3]])
    tmp1['l2'] = np.array([[ 0,  4],
                            [ 4,  5],
                            [57,  3]])
    tmp1['t3x'] = np.array([[188, 401, 400],
                           [188, 187, 401],
                           [187, 205, 401],
                           [324, 306, 417],
                           [306, 305, 417],
                           [305, 416, 417]])

    # replace values
    print(tmp1.strip(1))

    # replace values with a callable function that gets the size of the array
    print(tmp1.strip(lambda x: [None]*x))

    return 0

if __name__ == "__main__":
    sys.exit(main())

This code outputs:

{'l2': 1, 'n1': 1, 't3x': 1}
{'l2': [None, None, None], 'n1': [None], 't3x': [None, None, None, None, None, None]}

It's all good, it works. But I'm not sure the part that tackles the callable function is the most efficient one, as I am creating the copy of the dictionary first, and then I'm replacing the values with another loop. Can this be done in a single pass?

Answer 1

Although it's often EAFP, nonetheless it's still sometimes preferable to LBYL.

In this case, by making sure value is always something callable, the rest becomes simple if you use a dictionary comprehension to create the value returned.

class CustomDict(dict):

    def strip(self, value=None):
        if not callable(value):
            saved_value, value = value, lambda _: saved_value
        return {k: value(v) for k,v in self.items()}

tmp1 = CustomDict()

tmp1['n1'] = np.array([[3]])
tmp1['l2'] = np.array([[ 0,  4],
                       [ 4,  5],
                       [57,  3]])
tmp1['t3x'] = np.array([[188, 401, 400],
                        [188, 187, 401],
                        [187, 205, 401],
                        [324, 306, 417],
                        [306, 305, 417],
                        [305, 416, 417]])

print(tmp1.strip(1))
print(tmp1.strip(lambda x: [None]*len(x)))

Answer 2

Why an inherited class?

I don't see anything about this problem that requires an inherited class. strip could just as easily be a free function:

def strip(dct, value=None):
    ...

Also strip isn't a particularly good name - it doesn't really resemble the str.strip() function. I don't know what a better one would be though:

Callable

If you're going to allow passing an arbitrary function, then I think you may as well pass the full value instead of just the length. If the caller wants the length, the caller can use len at that end. This'll enhance all the things you can do.

Also, in Python, EAFP: it's easier to ask for forgiveness than permission. Just try to use the callable as a callable and catch the TypeError:

def strip(dct, value=None):
    try:
        tmp = dict.fromkeys(dct.keys())
        for k,v in dct.items():
            tmp[k] = value(v)
        return tmp
    except TypeError:
        return dict.fromkeys(dct.keys(), value)

Dictionary comprehension

For the except case - fromkeys() makes sense. But in the callable case, not really. You're creating a new dictionary and then re-modifying it. Better to just create the correct dictionary up front:

def strip(dct, value=None):
    try:
        return {k: value(v)
            for k,v in dct.items()
            }
    except TypeError:
        return dict.fromkeys(dct.keys(), value)