Python Add Dataclass to Set

Viewed 104

I have a dataclass defined like so:

from typing import List
from dataclasses import dataclass, field

@dataclass
class Speaker:
    id: int
    name: str
    statements: List[str] = field(default_factory=list)

    def __eq__(self, other):
        return self.id == other.id and self.name == other.name

    def __hash__(self):
        return hash((self.id, self.name))

and I have a list of names and statements which I want to combine. Each item in the list is going to have an id which may be shared any number of times. I want to append the statement part of each item in the list to the set of speakers.

This is what I have so far:

test = [(1, 'john', 'foo'),(1, 'john', 'bar'),(2, 'jane', 'near'),(2, 'george', 'far')]

speakers = set()
for i in test:
    id, name, statement = i
    Speaker(id, name)
    
    # This line needs to change
    speakers.add(Speaker(id, name, [statement]))

print(speakers)

Current output

{Speaker(id=1, name='john', statements=['foo']), Speaker(id=2, name= 'jane', statements=['near']), Speaker(id=2, name= 'george', statements=['far'])}

What I want

{Speaker(id=1, name='john', statements=['foo', 'bar']), Speaker(id=2, name= 'jane', statements=['near']), Speaker(id=2, name= 'george', statements=['far'])}

Please let me know if you have any suggestions. The number fields are subject to change (I may add title, ect.), so converting to a dict probably won't work.

Edit: Added an extra field called name to clarify the situation.

3 Answers

One way to do this is like so:

for i in test:
    id, statement = i
    new_speaker = Speaker(id, [statement])
    
    match = next((x for x in speakers if x == new_speaker), None)
    if match:
        match.statements.append(statement)
    else:
        speakers.add(new_speaker)

Not sure if there is a way which can do so in less time though, please comment if you think of any.

Instead of using a set, I think it makes more sense to use a dict type in this case. This should be O(1) for lookup, and also this way we can avoid next to find an element by its Speaker hashed value. Example below.

from pprint import pprint
from typing import List
from dataclasses import dataclass, field


@dataclass
class Speaker:
    id: int
    name: str
    statements: List[str] = field(default_factory=list)

    def __eq__(self, other):
        return self.id == other.id and self.name == other.name

    def __hash__(self):
        return hash((self.id, self.name))


test = [(1, 'john', 'foo'), (1, 'john', 'bar'), (2, 'jane', 'near'), (2, 'george', 'far')]
speakers = {}

for id, name, statement in test:
    key = Speaker(id, name)
    speaker = speakers.setdefault(key, key)
    speaker.statements.append(statement)

print(speakers)
print()
pprint(list(speakers.values()))

Out:

{Speaker(id=1, name='john', statements=['foo', 'bar']): Speaker(id=1, name='john', statements=['foo', 'bar']), Speaker(id=2, name='jane', statements=['near']): Speaker(id=2, name='jane', statements=['near']), Speaker(id=2, name='george', statements=['far']): Speaker(id=2, name='george', statements=['far'])}

[Speaker(id=1, name='john', statements=['foo', 'bar']),
 Speaker(id=2, name='jane', statements=['near']),
 Speaker(id=2, name='george', statements=['far'])]

Edit: as the dataclass itself is hashable, here I think it makes sense to store the Speaker object itself (uniquely identified by id, name, any other fields that are defined in __hash__) as both the key and value. This is a little roundabout, so if you wanted you could also store the tuple of the values you want to hash - i.e. (id, name) - which should also work. In either case, this should still be more efficient since it uses dict.setdefault which is still O(1) time to best of my knowledge.

If a set covers all use cases you have, but for merging one or more fields, you can customize a set, so that all the set-changing calls do merge the attributes.

The easiest way to do that is to write a class derived from collections.abc.MutableSet, which will both minimize the number of methods to be implemented, and ensure all changing calls go through your code (unlike, for example, subclassing Python's native set).

On a second thought: havign internally a set won't cut it - a mapping is needed so that it is easy to retrieve an specific element in order to update it. This approach allows us to keep a set interface, though, and we may add "update" and "get" methods to it:



from collections.abc import MutableMapping

class MergingSet(MutableSet):
    def __init__(self):
        # do not allow an initial content to avoid corner cases.
        # call .update  later to feed several instances at once
        self.data = dict()
        
    def __contains__(self, item):
        return item in self.data
    
    def __iter__(self):
        return iter(self.data)
    
    def __len__(self):
        return len(self.data)
    
    def discard(self, item):
        del self.data[item]
        self.data.discard(item)
        
    def add(self, item):
        if item in self:
            our_item = self.data[item]
            # hardcoded attribute merging:
            our_item.statements.extend(item.statements)
        else:
            self.data[item] = item
    
    def get(self, item):
        return self.data[item]  # returns the matching instance in the inner data, with the mutable fields
                                # modified.
                                
    def update(self, iterable):
        for item in iterable:
            self.add(item)

    def get(self, item, default=None):
        return self.data.get(item, default)    
            
    def __repr__(self):
        return repr(set(self.data.keys()))

And running this with your example:



In [28]:     
    ...: test = [(1, 'john', 'foo'),(1, 'john', 'bar'),(2, 'jane', 'near'),(2, 'george', 'far')]
    ...: 

In [29]: speakers = MergingSet()

In [30]: for i in test:
    ...:     id, name, statement = i
    ...:     Speaker(id, name)
    ...:     
    ...:     # This line needs to change
    ...:     speakers.add(Speaker(id, name, [statement]))
    ...:     

In [31]: speakers
Out[31]: {Speaker(id=2, name='jane', statements=['near']), Speaker(id=1, name='john', statements=['foo', 'bar']), Speaker(id=2, name='george', statements=['far'])}

Related