Grouper des objets selon un certain critère en python

Ajourd’hui, on fait des catégories. J’ai eu besoin d’établir des listes d’objets selon un certain critère, et retrouver mes listes par la valeur du critère. Par exemple, dans une liste d’objet, grouper les fruits, les animaux, et les machines-outils.

Imaginons que j’ai une liste d’objets qui portent tous le champ type.

Je pourrais procéder via une méthode évidente comme suit.

objects = something_to_get_objects()
result = dict()
for o in objects:
    if o.type not in result:
        result[o.type] = list()
    result[o.type].append(o)

Au lieu d’utiliser un champ de l’objet, je pourrais tout autant utiliser une fonction qui me retournerait une clé pour cet objet.

Ça fonctionne, mais je trouve que ça pourrait être plus élégant. Je pourrais par exemple déjà utiliser un defaultdict pour que chaque nouvelle clé à laquelle j’accéderai soit déjà peuplée par une liste vide.

from collections import defaultdict
objects = something_to_get_objects()
result = defaultdict(list)
for o in objects:
    result[o.type].append(o)

C’est déjà plus propre, mais j’ai l’impression que la lib itertools me donnerait bien des optimisations pour un très grand nombre d’objets.

Et voilà comment faire.

objects = something_to_get_objects()
# Sort is necessary for the next step
sorted_objects = sorted(objects, key=lambda x: x.type)

# groupby tool return a groupby object, which iterator over grouper objects
groupby = itertools.groupby(sorted_objects, key=lambda x: x.type)

# grouper objects are 2-tuples with a key and an iterator
# x[0] is the key
# x[1] is the iterator and I evaluate it to a list
result = dict((x[0], list(x[1]) for x in groupby)

# It's a dict!
for o in result['power-tool']: print o

J’aurais pu tout écrire sur une seule ligne mais ça aurait été overkill, non? Pour une fois, hélas, le code qui utilise itertools me paraît moins compact et élégant que la boucle avec le defaultdict. Est-il par contre plus performant, surtout pour de grosses listes en entrée?

Je suis par ailleurs obligé d’évaluer l’itérateur du grouper au moment où je construis le dictionnaire, ce que je trouve dommage. Mais ça s’explique par le fait qu’une fois l’itération sur les clés faite, l’itération sur les valeurs de la clé précédente n’est plus possible (le code de groupby est fait comme ça).

S’il s’agissait par contre d’utiliser ça pour itérer à la volée, ça serait parfait! Maintenant, à vous de voir quelle méthode vous convient le mieux.

Pour voir ou revoir le tri des listes, c’est par ici.