CODEWITHSUNDEEP
pythonpython-typing
Adam Lombard
Adam Lombard

Reputation: 334

Python type annotation for list of potential return values

Can a Python function be annotated with a list of potential return values?

RAINY = 'rainy'
SUNNY = 'sunny'
CLOUDY = 'cloudy'
SNOWY = 'snowy'
TYPES_OF_WEATHER = [RAINY, SUNNY, CLOUDY, SNOWY]

def today_is_probably(season) -> OneOf[TYPES_OF_WEATHER]:
  if season is 'spring':
    return RAINY
  if season is 'summer':
    return SUNNY
  if season is 'autumn':
    return CLOUDY
  if season is 'winter':
    return SNOWY

Upvotes: 2

Views: 610

Answers (1)

michaeldel
michaeldel

Reputation: 2385

The Literal type seems to be what you are looking for here:

from typing import Literal

def today_is_probably(season) -> Literal['rainy', 'sunny', 'cloudy', 'snowy']:
    ...

You may also consider defining an enum and use it as the type. It seems to be the most appropriate way here

import enum

class Weather(enum.Enum):
    RAINY = 'rainy'
    SUNNY = 'sunny'
    CLOUDY = 'cloudy'
    SNOWY = 'snowy'

def today_is_probably(season) -> Weather:
    ...

Upvotes: 4

Related Questions