Home > Blockchain >  Python Add Dataclass to Set
Python Add Dataclass to Set

Time:06-18

I have a dataclass defined like so:

from typing import List
from dataclasses import dataclass, field

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

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

    def __hash__(self):
        return self.id

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, 'foo'),(1,'bar'),(2,'near'),(2,'far')]

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

print(speakers)

Current output

{Speaker(id=1, statements=['foo']), Speaker(id=2, statements=['near'])}

What I want

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

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

CodePudding user response:

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.

CodePudding user response:

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 id property. Example below:

from typing import List
from dataclasses import dataclass, field


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

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


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

for id, statement in test:
    speaker = speakers.get(id)
    if speaker is None:
        speaker = speakers[id] = Speaker(id)

    speaker.statements.append(statement)

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

Out:

{1: Speaker(id=1, statements=['foo', 'bar']), 2: Speaker(id=2, statements=['near', 'far'])}
[Speaker(id=1, statements=['foo', 'bar']), Speaker(id=2, statements=['near', 'far'])]
  • Related