Strategies for Optional attributes that become non-Optional during runtime?

[jeremyn]

You might have a class with an attribute that you want to accept None early in its life, but after some specific point it's guaranteed to be non-None. For example:

from dataclasses import dataclass
from typing import Optional

@dataclass
class MyClass:
    val: str
    key: Optional[int] = None

def get_objs():
    objs = [MyClass("b"), MyClass("a")]
    for index, obj in enumerate(sorted(objs, key=lambda o: o.val)):
        obj.key = index
    return objs

objs = get_objs()
print(objs) # [MyClass(val='b', key=1), MyClass(val='a', key=0)]
objs[0].key + 1 # typing error: can't add to None

In this case anything using objs is assured of having non-None keys, but is forced to do if obj.key is not None checks everywhere to avoid typing complaints.

There are workarounds, such as:

• using some nonsensical default like -1 for key
• not defining key during initialization
• defining a shadow _key Optional[int] attribute and presenting key as an int property
• rewriting get_objs such that MyClass objs are only instantiated at the end when key is known

These are all sort of hacking around the type system though and may not be available or have other problems.

Any thoughts or general strategies on dealing with this? Thanks.

[JelleZijlstra]

I don't think the type system provides any direct way to do what you want here, and it's hard to imagine a clean way to add something like this. Most type checkers should support something like assert obj.key is not None, but that gets repetitive.

Personally I would try to rewrite this code so that the key attribute is known when the object is initialized.

[jeremyn]

@JelleZijlstra @srittau @sobolevn Thanks for the responses. I'm marking this one as the answer only because it was first and I can't mark more than one.

[srittau]

I think this is hard to do in the current type system. Usually I try to follow @JelleZijlstra's advice, since I'm a big fan of immutable object. But if that's not reasonable to do, I would probably define MyClass.key as non-optional and deal with the optional fields inside get_objs(), using # type: ignore, casts, or a second "constructor" type. This way, the type problems are localized to a method that "owns" the construction and the rest of your system is free from type problems.

I also like the solution that you hinted at, where you define a private _key field, for some situations:

class MyClass:
    def __init__(self, val: str, key: int | None = None) -> None:
        self._key = key

    @property
    def key(self) -> int:
        if self._key is None:
            raise RuntimeError("key not set")
        return self._key

    @key.setter
    def key(self, key: int) -> None:
        self._key = key

If there isn't already a custom property that does something like that in the standard library, it should be fairly easy to create one to avoid most of the boilerplate.

[sobolevn]

It is not really a good idea to use mutable structures and typing together. Especially in cases when you change your type: in this case you are changing Optional[int] to be just int.

But, you can still express this with Python's type system and three hacks:

from dataclasses import dataclass
from typing import Optional, Generic, TypeVar, List, cast

_PayloadType = TypeVar('_PayloadType')

@dataclass
class _MyClass(Generic[_PayloadType]):
    val: str
    # NOTE: hack1
    key: _PayloadType = None  # type: ignore

# NOTE: hack2
MyClass = _MyClass[Optional[int]]
MyFullClass = _MyClass[int]

def get_objs() -> List[MyFullClass]:
    objs = [MyClass("b"), MyClass("a")]
    for index, obj in enumerate(sorted(objs, key=lambda o: o.val)):
        obj.key = index
    # NOTE: hack3
    return cast(List[MyFullClass], objs)

objs = get_objs()
objs[0].key + 1 # ok

Hacks:

1. We use # type: ignore to make sure that we can really assign None to _PayloadType
2. We create two invariants of _MyClass type: MyClass that has _PayloadType of Optional[int] and MyFullClass with just int. You can even use NewType here in some cases
3. We use an explicit cast to tell the type checker that we have done some work it cannot understand

It works 🎉

But I am not really happy with this code. Here's an alternative:

from dataclasses import dataclass
from typing import Optional, List

@dataclass
class MyClass:
    val: str

@dataclass
class MyEmptyClass(MyClass):
    key: Optional[int] = None

@dataclass
class MyFullClass(MyClass):
    key: int

def get_objs() -> List[MyFullClass]:
    objs = [MyEmptyClass("b"), MyEmptyClass("a")]
    new_objs = []
    for index, obj in enumerate(sorted(objs, key=lambda o: o.val)):
        new_objs.append(MyFullClass(val=obj.val, key=index))
    return new_objs

objs = get_objs()
objs[0].key + 1 # ok

It is cleaner, but uses two lists and copies extra data. It is all about trade-offs.

[jeremyn]

@sobolevn I like your second alternative but I think this is a slight improvement:

@dataclass
class _MyClass:
    val: str
    key: Optional[int] = None

class MyClass(_MyClass):
    key: int

def get_objs() -> list[MyClass]:
    objs = [_MyClass("b"), _MyClass("a")]
    for index, obj in enumerate(sorted(_objs, key=lambda o: o.val)):
        obj.key = index
    new_objs = cast(list[MyClass], objs) # or just return the cast directly
    return new_objs

We have one less class to worry about, and also objs and new_objs are literally the same list -- the ids are the same -- so there's no extra data or resources being used. What do you think?

[sobolevn]

I like it! 👍

[jeremyn]

I agree with everyone that letting the attribute just be key: int and then dealing with it in the code instead of working around the type system, feels more correct.

However, perhaps a feature request to the type system would be to allow for attribute-level casts, such as cast(obj, "key", int), to (as @srittau said in another context) avoid some boilerplate. Does anyone think this would be a good idea? I'm not really sure myself.

[jeremyn]

Or (better) cast(int, obj, attr="key").

Of course, to illustrate a problem with this idea, the actual example I'm thinking about is nested: the key is part of another class, also used elsewhere, that's composed into the main class. Something like this:

class KeyClass:
    attr1: str
    attr2: Optional[int]

class MyClass:
    val: str
    key: KeyClass

So with attribute-level casting, in this case, I'd want to override the type for MyClass.key.attr2, and only for those KeyClasses attached to MyClasses, not all KeyClasses everywhere. I'm not sure of a clean way to request that, and I'm sure implementation would be much worse.

[JelleZijlstra]

Just so I understand the idea, the way this cast would work is that if I do obj2 = cast(int, obj, attr="key"), obj2 would have the same type as obj except that obj.key is now an int?

[sobolevn]

I'm not sure of a clean way to request that, and I'm sure implementation would be much worse.

Yes, I agree. It would be a nightmare to implement. We cannot touch TypeInfo (where all type-specific information lives in mypy), because it would affect other instances. So, it would have to be some kind of an overwriting hack. Big "no" from me.

[jeremyn]

Just so I understand the idea, the way this cast would work is that if I do obj2 = cast(int, obj, attr="key"), obj2 would have the same type as obj except that obj.key is now an int?

Yes, that would be the outcome. Another way to think of it would be like how the patch decorator in unittest.mock works: you would "patch" the type for some object's attribute and the type checker would work with the patched type instead of the original one for that object going forward. Conceptually behind the scenes I expect it could create anonymous wrapper types, like I did in my reply to @sobolevn. Even nested types shouldn't be too bad.

I don't know the details of how this would be implemented but it doesn't seem impossible, just tedious. Also, from a user perspective, patch isn't known for being super intuitive. However the biggest argument against it might be that it would encourage bad typing habits, people just patching the type to get things to pass instead of organizing their code better.

[tkdrob]

I've been pursuing making data classes that satisfy mypy for external packages without having to cast them. When using @dataclass(init=False) to decorate a class, I noticed that pylance does not care about the order of attributes with a default and those without one. But mypy does care. So I wonder if this is by design or simply overlooked.

[erictraut]

I think pyright/pylance is correct here. The order does not matter when init=False.

# Does not generate a runtime exception
@dataclass(init=False)
class DC1:
    attr1: int = 1
    attr2: int

# Generates a runtime exception
@dataclass
class DC2:
    attr1: int = 1
    attr2: int

Mypy generates errors in both of these cases. Pyright generates an error only for the second case, which matches the runtime behavior.

While investigating this, I discovered an edge case that pyright wasn't handling correctly. If init=True but there is an existing __init__ method, pyright currently assumes that the ordering doesn't matter because no __init__ method is synthesized. However, it appears that the runtime still raises an exception in this situation. I've updated pyright's logic to match the runtime in this case.

# Generates a runtime exception
@dataclass
class DC2:
    attr1: int = 1
    attr2: int

    def __init__(self):
        pass