disjoint unions in mypy?

[kurtgn]

I am trying to customize a function that takes an object of a union type and accesses different attributes depending on the type. The type of the object is specified by a type attribute.

from typing import Union

from dataclasses import dataclass


TEXT = 'TEXT'
IMAGE = 'IMAGE'


@dataclass
class TextMessage:
    text: str
    type = TEXT


@dataclass
class ImageMessage:
    url: str
    type = IMAGE


def output(message: Union[TextMessage, ImageMessage]) -> None:
    if message.type == TEXT:
        print(message.text)

Mypy does not like this, saying:

 error: Item "ImageMessage" of "Union[TextMessage, ImageMessage]" has no attribute "text"

And I understand it. I am trying to guess one attribute based on the value of another one. This is poor engineering.

However, the same thing works in other languages, for example, in JS Flow:

// @flow

type TextMessage = {
	type: 'TEXT',
  	text: string
};
type ImageMessage = {
	type: 'IMAGE',
  	url: string
};

function output(message: TextMessage | ImageMessage) {
  if (message.type === 'TEXT') {
    console.log(message.text);
  }
} 

Live Flow link: here

So the question is, is there a way to make mypy understand my conditions? Maybe using Literal types? I tried to figure this out but the docs are a bit confusing.

P.S. I can always do if isinstance(message, TextMessage) and this will satisfy mypy, but putting lots of isinstance() in my code seems even more hacky.

[ethanhs]

P.S. I can always do if isinstance(message, TextMessage) and this will satisfy mypy, but putting lots of isinstance() in my code seems even more hacky.

Aren't you essentially trading off putinig lots of isinstance(message, TextMessage) with lots of if message.type == 'TEXT'? The first is 1) more Pythonic 2) better supported 3) more readable and 4) doesn't require you to add a type tag to all of your classes.

As for using Literal to solve your issues, under "Extra" features to add in #5935 there is "Type narrowing with equality and inequality checks", which is what you want.

[ilevkivskyi]

But why this

def output(message: Union[TextMessage, ImageMessage]) -> None:
    if message.type == TEXT:
        print(message.text)

is better than this

def output(message: Union[TextMessage, ImageMessage]) -> None:
    if isinstance(message, TextMessage):
        print(message.text)

?

Anyway, I propose to close this since this is essentially the same as "Type narrowing with equality and inequality checks" item from #5935 as @ethanhs mentioned.

[Garrett-R]

Anyway, I propose to close this

Agreed!

Interestingly, TypeScript does offer this feature (discriminating unions), but @mneira10 pointed out that they're more important in TS given that TS doesn't have a clean way of doing isinstance like Python does. You could write user-defined type guards (example), but that's messy.

[JukkaL]

Besides, the original example works already if using Final and literal types:

from typing import Final, Literal
from dataclasses import dataclass


TEXT: Final = 'TEXT'
IMAGE: Final = 'IMAGE'


@dataclass
class TextMessage:
    type: Literal['TEXT']
    text: str


@dataclass
class ImageMessage:
    type: Literal['IMAGE']
    url: str


def output(message: TextMessage | ImageMessage) -> None:
    if message.type == TEXT:
        print(message.text)  # No error