CODEWITHSUNDEEP
pythonenums
Fiebbo
Fiebbo

Reputation: 33

elegant way of counting items in a list by enum value in python

I have a list of objects where every object contains a type which is an IntEnum value. I'm looking for an elegant way to count the number of objects in my list by their type(IntEnum) value.

Below is a fully functioning snippet of what I came up with but it seems clunky to me. For instance adding a new type to my enum would require manual modification to other parts of the code.

from dataclasses import dataclass
from enum import IntEnum, unique
from random import randint
from textwrap import dedent
from collections import Counter

@unique
class ObjType(IntEnum):
    OT_START = 0
    OBJTYPE_0 = 1
    OBJTYPE_1 = 2
    OBJTYPE_2 = 3
    OT_END = 4
    # many more types...

@dataclass
class ObjListAttributes:
    OBJTYPE_0_count = 0
    OBJTYPE_1_count = 0
    OBJTYPE_2_count = 0

@dataclass
class ObjClass:
    t: ObjType
    attr_1: int = 0
    attr_2: float = 0
    # many more attributes...

# create list of objects with random types
objList = [ObjClass(randint(ObjType.OT_START + 1, ObjType.OT_END - 1)) for x in range(100)]

# count number of objects in list by type **THIS IS THE CLUNKY PART**
listAttrs = ObjListAttributes()
for x in objList:
    if x.t == ObjType.OBJTYPE_0:
        listAttrs.OBJTYPE_0_count += 1
    elif x.t == ObjType.OBJTYPE_1:
        listAttrs.OBJTYPE_1_count += 1
    elif x.t == ObjType.OBJTYPE_2:
        listAttrs.OBJTYPE_2_count += 1

print(dedent(f'''\
    list contains:
        {listAttrs.OBJTYPE_0_count} {ObjType.OBJTYPE_0.name} objects
        {listAttrs.OBJTYPE_1_count} {ObjType.OBJTYPE_1.name} objects
        {listAttrs.OBJTYPE_2_count} {ObjType.OBJTYPE_2.name} objects
    '''))

To give you some context, the snippet is part of some code that unpacks and analyzes binary data that was generated in C code. Therefore, using enums felt natural for solving the problem. I'm open to using other data structures if they result in more elegant code.

Upvotes: 1

Views: 587

Answers (2)

9769953
9769953

Reputation: 12201

A bit hacky, but you could implement a __hash__ method to your ObjClass class that just uses t:

@dataclass
class ObjClass:
    t: ObjType
    attr_1: int = 0
    attr_2: float = 0
    # many more attributes...

    def __hash__(self):
        return hash(self.t)

Then,

counter = Counter(objList)

will work.

Afterwards, you can reset the __hash__ method for the ObjClass:

ObjClass.__hash__ = None

It is perhaps unfortunate that Counter doesn't have a key parameter as for sort and the like. Then again, since ObjClass isn't hashable, its instances can't function as a key for a Counter instance (or any dict) either, so it wouldn't help in this case.

Upvotes: 1

kubatucka
kubatucka

Reputation: 565

Check out Counter. But beware using it requires hashable dict keys so you can't just use the class object.

from collections import Counter

c= Counter(map(lambda x:x.t,objList))  # it should be x.t.value if t is an intEnum

Upvotes: 3

Related Questions