Search code examples
pythonpython-typing

Check if a field is typing.Optional


What is the best way to check if a field from a class is typing.Optional?

Example code:

from typing import Optional
import re
from dataclasses import dataclass, fields

@dataclass(frozen=True)
class TestClass:
    required_field_1: str
    required_field_2: int
    optional_field: Optional[str]

def get_all_optional_fields(fields) -> list:
    return [field.name for field in fields if __is_optional_field(field)]

def __is_optional_field(field) -> bool:
    regex = '^typing.Union\[.*, NoneType\]$'
    return re.match(regex, str(field.type)) is not None

print(get_all_optional_fields(fields(TestClass)))

Where fields is from dataclasses, I wanna list all the Optional fields. What I'm doing at this moment to solve it, is using a Regex-based on the field name, but I don't like this approach. Is there a better way of doing it?


Solution

  • Optional[X] is equivalent to Union[X, None]. So you could do,

    import re
    from typing import Optional
    
    from dataclasses import dataclass, fields
    
    
    @dataclass(frozen=True)
    class TestClass:
        required_field_1: str
        required_field_2: int
        optional_field: Optional[str]
    
    
    def get_optional_fields(klass):
        class_fields = fields(klass)
        for field in class_fields:
            if (
                hasattr(field.type, "__args__")
                and len(field.type.__args__) == 2
                and field.type.__args__[-1] is type(None)
            ):
                # Check if exactly two arguments exists and one of them are None type
                yield field.name
    
    
    print(list(get_optional_fields(TestClass)))