CODEWITHSUNDEEP
pythonpython-typing
Martin O Leary
Martin O Leary

Reputation: 691

How can I get the generic type of a parameterized generic used to annotate a dataclass property?

I have a dataclass:

import dataclasses

@dataclasses.dataclass
class Foo:
    first: dict[str, int]

How can I determine that the generic type dict is the type expected with a str to int mapping? I want to be able to do something like:

for field in dataclasses.fields(Foo):
    if generic_type_of_field == dict:
        <<do something specific to this field needing to be a dict>>
    elif generic_type_of_field == list:
        <<do something specific to this field needing to be a list>>

I have tried using dataclasses.fields(Foo) to get the fields, then taking field.type. This just gives me dict[str, int] which is a <class 'types.GenericAlias'> type. Is there something more I can do here?

I can get the paramaterizations of the dict generic via get_args(field.type) which gives me (<class 'str'>, <class 'int'>) (which I do need) but I don't have run out of knowledge on even searching on the topic now so I am asking here.

Upvotes: 1

Views: 205

Answers (1)

Daniil Fajnberg
Daniil Fajnberg

Reputation: 18568

If you want to be able to tell the original generic type of a specified generic type, you can use typing.get_origin.

The problem is that it will return None, if you pass it a non-specified generic type. If you want to account for that, the simplest solution I can think of something along these lines:

import dataclasses
from typing import get_origin


@dataclasses.dataclass
class Foo:
    first: dict[str, int]
    second: dict
    third: list[float]


def generic_base(t: type) -> type:
    origin = get_origin(t)
    if origin is not None:
        return origin
    return t


if __name__ == '__main__':
    for field in dataclasses.fields(Foo):
        if generic_base(field.type) is dict:
            print(field.name, "is a dict")
        elif generic_base(field.type) is list:
            print(field.name, "is a list")

The output:

first is a dict
second is a dict
third is a list

Upvotes: 1

Related Questions