简体   繁体   中英

Python dataclass to implement choice type

I have a set of dataclasses say:

    from dataclasses import dataclass, asdict, InitVar

    @dataclass
    class Item:  
        name: str = None
        identifier: int = None
        
    @dataclass
    class Container:
        item: Item
        cid: int
        cname: str = None
    

When I do:

c = Container(Item(name="item-1"), cid=10)
asdict(c)

I get:

{'item': {'name': 'item-1', 'identifier': None},'cid': 10,'cname': None}

But in my schema Item is a "choice" type so I only want to include "name" or "identifier" in "asdict" depending on which of those are actually set (ONLY for Item type).

something like:

{'item': {'name': 'item-1'},'cid': 10,'cname': None}

OR:

{'item': {'identifier': 'id-1'},'cid': 10,'cname': None}

My original code is much more complex and the relationships are more nested so I'm looking for a solution which I can apply to the specific dataclass. I tried manipulating the __dict__ to add attributes in __post_init__ but that didn't work. For eg I tried

from dataclasses import dataclass, asdict, InitVar

@dataclass
class Item:  
    name: InitVar[str] = None
    identifier: InitVar[int] = None
        
    def __post_init__(self, name, identifier):
        if name:
            self.name = name
        elif identifier:
            self.identifier = identifier
        print(self.__dict__)
        
    
@dataclass
class Container:
    item: Item
    identifier: int
    cname: str = None
    

c = Container(Item(name="item-1"), cid=10)
asdict(c)

but that prints

{'item': {}, 'cid': 10, 'cname': None}

It is probably not what you want, but at this time the only way forward when you want a customized dict representation of a dataclass is to write your own .asdict method.

Here's a suggested starting point (will probably need tweaking):

from dataclasses import dataclass, asdict


@dataclass
class DataclassAsDictMixin:
    def asdict(self):
        d = asdict(self)
        for field, value in ((f,v) for f,v in vars(self).items() if f in d):
            try:
                value = value.asdict()
            except AttributeError:
                pass
            else:
                d.update([(field, value)])
        return d


@dataclass
class Item:  
    name: str = None
    identifier: int = None

    def asdict(self):
        d = asdict(self)
        for k,v in d.copy().items():
            if v is None:
                del d[k]
        return d

@dataclass
class Container(DataclassAsDictMixin):
    item: Item
    cid: int
    cname: str = None


if __name__ == "__main__":
    c1 = Container(Item(name="item-1"), cid=10)
    assert c1.asdict() == {'item': {'name': 'item-1'}, 'cid': 10, 'cname': None}
    c2 = Container(Item(identifier="id-1"), cid=10)
    assert c2.asdict() == {'item': {'identifier': 'id-1'}, 'cid': 10, 'cname': None}

I haven't tested it with a list of values, but try this approach:

from copy import copy
from dataclasses import dataclass, fields

from validated_dc import ValidatedDC


@dataclass
class Base(ValidatedDC):

    def is_correct_value(self, value):

        return True

    def as_dict(self):

        result = {}
        nested = tuple(self.get_nested_validated_dc())

        for field in fields(self):
            value = copy(getattr(self, field.name))

            if isinstance(value, list):
                for item in value:
                    if isinstance(item, nested):
                        item = item.as_dict()

            if isinstance(value, nested):
                value = value.as_dict()

            if self.is_correct_value(value):
                result[field.name] = value

        return result


@dataclass
class ItemBase(Base):

    def is_correct_value(self, value):

        return False if value is None else True


@dataclass
class Item(ItemBase):

    name: str = None
    identifier: int = None


@dataclass
class Container(Base):

    item: Item
    cid: int
    cname: str = None


c = Container(Item(name="item-1"), cid=10)
assert c.as_dict() == {'item': {'name': 'item-1'}, 'cid': 10, 'cname': None}

ValidatedDC: https://github.com/EvgeniyBurdin/validated_dc

The technical post webpages of this site follow the CC BY-SA 4.0 protocol. If you need to reprint, please indicate the site URL or the original address.Any question please contact:yoyou2525@163.com.

 
粤ICP备18138465号  © 2020-2024 STACKOOM.COM