typing module - String Literal Type [duplicate]

Question

This question already has answers here:

Type hint for a function that returns only a specific set of values (3 answers)

I'm using the new Python 3.5 module typing and it has been joyous.

I was wondering how one might specify a type based on an exact string literal. For example, a function is guaranteed to return one of the four strings - "North", "West", "East", "South - how can we express that as a specific type variable, instead of just str.

I looked through the documentation, finding the Union type and the TypeVar function, but was unable to find an answer.

An example function expressing this problem:

def compute_quadrant(x: int, y: int) -> str:
    if x > 0 and y > 0:
        return 'I'
    elif x < 0 and y > 0:
        return 'II'
    elif x < 0 and y < 0:
        return 'III'
    elif x > 0 and y < 0:
        return 'IV'

Instead of just returning str, I'd like to to return a more specific type that is one of four values - "I", "II", "III", or "IV".

In Typescript, one can do: type Quadrant = "I" | "II" | "III" | "IV" - is there any nice Python sugar for this use case with the typing module?

Answer 1

Disregarding the typing module you asked about, one solution to your problem could be the use of an Enum as supposed in multiple comments. The code for this would look like this:

from enum import Enum

class Quadrant(Enum):
    I = 1
    II = 2
    III = 3
    IV = 4

def compute_quadrant(x: int, y: int) -> Quadrant:
    if x > 0 and y > 0:
        return Quadrant.I
    elif x < 0 and y > 0:
        return Quadrant.II
    elif x < 0 and y < 0:
        return Quadrant.III
    elif x > 0 and y < 0:
        return Quadrant.IV
    # return None  # this is what happens without an else clause!

if __name__ == "__main__":
    quad = compute_quadrant(1, -1)
    print(quad, type(quad))              # -> Quadrant.IV <enum 'Quadrant'>
    print(quad.name, type(quad.name))    # -> IV <class 'str'>
    print(quad.value, type(quad.value))  # -> 4 <class 'int'>

As you can see you can use the Enums name and value. The name is one of the strings you asked for.

One issue I see here is the missing else clause in the function and mypy's current behaviour of accepting None as a valid return value for Quadrant. This should be handled manually.