Useful Python functions

common.py

#!/usr/bin/env python
import collections

# DO NOT USE GENSHI'S FLATTEN - it is not a generator and will chew up CPU and RAM 
def flatten(l):
    """
    Flatten all iterables (lists, tuples, dicts, etc), into one lazy generator-stream
    """
    if isinstance(l, dict):
        l = l.items()
    for el in l:
        if isinstance(el, collections.Iterable) and not isinstance(el, basestring):
            for sub in flatten(el):
                yield sub
        else:
            yield el

def flattenCast(l, cast_to=list):
    """
    Flatten all iterables (lists, tuples, dicts, etc) and cast into a new collection.
    The default is to cast it into a list
    """
    return cast_to(flatten(l))

def selectKeys(d, keys):
    """
    Like Clojure's select-keys: sub-dict selection.
    Arguments:
        d - a dict, the source dictionary
        keys - a list, the keys you want to select from `d`, for your new dict
                You can select inner-dicts with lists:
                    Given: {"a": 1, "b": 2, "c": {"d": 4}}
                    keys=[["c", "d"]]       =>  {"d": 4"}
                    keys=["a", ["c", "d"]]  =>  {"a": 1, "d": 4}
                    keys=["a", "c"]         =>  {"a": 1, "c": {"d": 4}}
                    keys=["a", ["c"]]       =>  {"a": 1, "c": {"d": 4}}
    Notes:
        This is non-destructive
    """
    #return reduce(lambda x, y: x.update({y[0]:y[1]}) or x,
    #        map(None, keys, map(d.get, keys)), {})
    # If Python < 2.7
    #return dict((k, d[k]) for k in keys if k in d)
    
    # Simple top-key selection
    #return {k: d[k] for k in keys if k in d}
    ret_dict = dict()
    for k in keys:
        if isinstance(k, list):
            last_k = k[-1]
            value = reduce(lambda dd, k: dd.get(k, dd), k, d)
            ret_dict[last_k] = value
        else:
            if k in d:
                ret_dict[k] = d[k]
    return ret_dict

test_common.py

class TestCommon(TestCase):

    def test_flattenCast(self):
        """
        Verify we can flatten all lists, lists of lists, etc.
        """
        test_data = [([1, 2, 3, 4],     [1, 2, 3, 4]),
                     ([1, 2, [3]],      [1, 2, 3]),
                     ([1, [2, [3]]],    [1, 2, 3]),
                     ([[[[1]]]],        [1]),
                     ((1, 2, 3, 4),     [1, 2, 3, 4]),
                     ({"a": 1},         ["a", 1]),
                     ([1, ["hello"]],   [1, "hello"]),
                     ({"mykey": 1},     ["mykey", 1])]
        for test_list, expected in test_data:
            actual = common.flattenCast(test_list)
            assert actual == expected
        for test_list, expected in test_data:
            expected = tuple(expected)
            actual = common.flattenCast(test_list, cast_to=tuple)
            assert actual == expected

    def test_selectKeys(self):
        """
        Verify we can do sub-dict selections ala selectKeys
        """
        test_dict = {"a": 1, "b": 2, "c": 3, "d": {"e": 5, "f":6}}
        test_data = [(["a", "c"],           {"a": 1, "c": 3}),
                     (["a"],                {"a": 1,}),
                     ([["a"]],              {"a": 1,}),
                     ([],                   {}),
                     ([["d", "e"]],         {"e": 5}),
                     (["a", ["d", "e"]],    {"a": 1, "e": 5}),
                     (["a", "d"],           {"a": 1, "d": {"e": 5, "f":6}}),
                     (["a", ["d"]],         {"a": 1, "d": {"e": 5, "f":6}})]
        for test_keys, expected in test_data:
            actual = common.selectKeys(test_dict, test_keys)
            assert actual == expected