Shay
Shay

Reputation: 1499

how can I narrow dataclass annotations (i.e., how can I update type hints after handling default None in post_init)?

I have a dataclass that can take a keyword value, or, if no value is specified, infer a value from other attributes.

import dataclasses

@dataclasses.dataclass
class RelatedValues:
    primary: float
    _: dataclasses.KW_ONLY
    secondary: float | None = None

def __post_init__(self):
    if self.secondary is None:
        self.secondary = self.primary

This code works, but it leaves me stuck with float | None as the type hint for .secondary even though .secondary cannot possibly be None after __post_init__.

cast-ing self.secondary in __post_init__ doesn't work. This does:

NULL_FLOAT = float(int(uuid.uuid4())

@dataclasses.dataclass
class RelatedValues:
    primary: float
    _: dataclasses.KW_ONLY
    secondary: float = NULL_FLOAT

def __post_init__(self):
    if self.secondary == NULL_FLOAT:
        self.secondary = self.primary

But it feels distinctly non-Pythonic.

This also works:

@dataclasses.dataclass
class RelatedValues:
    primary: float
    _: dataclasses.KW_ONLY
    _secondary: float | None = None

def __post_init__(self):
    if self._secondary is None:
        self.secondary = self.primary
    else:
        self.secondary = self._secondary

or this:

@dataclasses.dataclass
class RelatedValues:
    primary: float
    _: dataclasses.KW_ONLY
    _secondary: float | None = None

@property
def secondary(self) -> float:
    if self._secondary is None:
        self.secondary = self.primary
    else:
        self.secondary = self._secondary

But the latter two are just mangling my kwargs for the sake of type narrowing, which kind of feels wrong.

What am I missing?

Upvotes: 0

Views: 139

Answers (1)

chepner
chepner

Reputation: 532208

If secondary will always be assigned a non-None value, it should not be typed as float | None in the first place. It's only the argument to __init__, used to initialize self.secondary, that might be None.

from dataclasses import dataclass, field, InitVar


@dataclass
class RelatedValues:
    primary: float
    secondary: float = field(init=False)
    secondary_: InitVar[float|None] = None

    def __post_init__(self, secondary_):
        if secondary_ is None:
            secondary_ = self.primary
        self.secondary = secondary_

These hoops are only necessary if you don't want to provide your own __init__ method in the first place. For example,

@dataclass(init=False)
class RelatedValues:
    primary: float
    secondary: float

    def __init__(self, primary: float, secondary: float = None):
        if secondary is None:
            secondary = primary
        self.primary = primary
        self.secondary = secondary

Upvotes: 1

Related Questions