Search code examples
pythonpython-3.xpython-dataclasses

How to remove dynamically fields from a dataclass?


I want to inherit my dataclass but remove some of its fields. How can I do that in runtime so I don't need to copy all of the members one by one?

Example:

from dataclasses import dataclass

@dataclass
class A:
    a: int
    b: int
    c: int
    d: int

@remove("c", "d")
class B(A):
    pass

Such that A would have a, b, c, d defined and B would only have a and b defined.


Solution

  • We can remove the particular fields from the __annotations__ dictionary as well as from __dataclass_fields__ and then rebuild our class using dataclasses.make_dataclass:

    def remove(*fields):
        def _(cls):
            fields_copy = copy.copy(cls.__dataclass_fields__)
            annotations_copy = copy.deepcopy(cls.__annotations__)
            for field in fields:
                del fields_copy[field]
                del annotations_copy[field]
            d_cls = dataclasses.make_dataclass(cls.__name__, annotations_copy)
            d_cls.__dataclass_fields__ = fields_copy
            return d_cls
        return _
    

    Note that we copy the annotations and the fields in order to not affect A. otherwise it would remove these fields from A as well and any other attempt to inherit A and remove again a field which we already removed would lead to an error. I.e:

    from dataclasses import dataclass
    
    @dataclass
    class A:
        a: int
        b: int
        c: int
        d: int
    
    @remove("c", "d")
    class B(A):
        pass
    
    @remove("b", "c")
    class C(A):
        pass
    

    This would give us a KeyError since we already removed "c" and this is no longer exists in A's dictionary.