CODEWITHSUNDEEP
pythonpython-typing
Alberto B
Alberto B

Reputation: 541

Check if a type is Union type in Python

I have defined a dataclass:

import dataclasses

@dataclasses.dataclass
class MyClass:
    attr1: int | None
    attr2: str | None

I can loop through the types of my attributes with:

for field in dataclasses.fields(MyClass):
    fieldname = field.name
    fieldtype = field.type

But how can I check if type 'str' is in 'fieldtype' or get the list of types inside the union type?

Upvotes: 1

Views: 1934

Answers (1)

user459872
user459872

Reputation: 24562

You can use __args__ attribute of union type.

>>> import dataclasses
>>> import types
>>> 
>>> @dataclasses.dataclass
... class MyClass:
...     attr1: int | None
...     attr2: str | None
... 
>>> for field in dataclasses.fields(MyClass):
...     fieldname = field.name
...     fieldtype = field.type 
...     has_str_exists = False
...     if isinstance(fieldtype, types.UnionType):
...         fieldtype = fieldtype.__args__
...         has_str_exists = str in fieldtype
...     print(f"{fieldname=} {fieldtype=} {has_str_exists=}")
... 
fieldname='attr1' fieldtype=(<class 'int'>, <class 'NoneType'>) has_str_exists=False
fieldname='attr2' fieldtype=(<class 'str'>, <class 'NoneType'>) has_str_exists=True

Upvotes: 2

Related Questions