Using defaultdict from the collections module instead of setdefault method of dict.

defaultdict gives a dictionary like object

This defaultdict object is kind of cool. Whenever a missing key is accessed, defaultdict will call the callable that was given to it (list in this case) and use the return value as the new value for that key.


In [6]:
from collections import defaultdict
def group_by(iterable, key_func=lambda x: x):
    groups = defaultdict(list)
    
    for item in iterable:
        value = groups[key_func(item)]
        print(value)
    return groups

In [7]:
for key, value in group_by([1, 2, 1, 2]).items():
    print(key, value)


[]
[]
[]
[]
1 []
2 []

In [8]:
def group_by(iterable, key_func=lambda x: x):
    groups = defaultdict(list)
    
    for item in iterable:
        groups[key_func(item)].append(item)
        
    return groups

In [9]:
for key, value in group_by([1, 2, 1, 2]).items():
    print(key, value)


1 [1, 1]
2 [2, 2]

In [11]:
from operator import itemgetter
import unittest

class GroupByTests(unittest.TestCase):

    """Tests for group_by."""

    def test_test_tuples_of_strings(self):
        animals = [
            ('agatha', 'dog'),
            ('kurt', 'cat'),
            ('margaret', 'mouse'),
            ('cory', 'cat'),
            ('mary', 'mouse'),
        ]
        animals_by_type = {
            'mouse': [('margaret', 'mouse'), ('mary', 'mouse')],
            'dog': [('agatha', 'dog')],
            'cat': [('kurt', 'cat'), ('cory', 'cat')],
        }
        output = group_by(animals, key_func=itemgetter(1))
        self.assertEqual(output, animals_by_type)

    def test_strings(self):
        words = ["Apple", "animal", "apple", "ANIMAL", "animal"]
        word_groups = {
            "apple": ["Apple", "apple"],
            "animal": ["animal", "ANIMAL", "animal"],
        }
        output = group_by(words, key_func=str.lower)
        self.assertEqual(output, word_groups)

    # To test the Bonus part of this exercise, comment out the following line
    # @unittest.expectedFailure
    def test_no_key_function(self):
        words = ["apple", "animal", "apple", "animal", "animal"]
        word_groups = {
            "apple": ["apple", "apple"],
            "animal": ["animal", "animal", "animal"],
        }
        output = group_by(words)
        self.assertEqual(output, word_groups)


if __name__ == "__main__":
    unittest.main(argv=['ignore-first-argument'], exit=False)


...
----------------------------------------------------------------------
Ran 3 tests in 0.002s

OK

In [ ]: