Python Using Literal type to restrict the parameter

There are some cases where a function requires a string as a parameter but wants to restrict the acceptable values. How can we implement it in Python?

Implementation without any restriction

Assume that there are three roles and we want to get the corresponding ID of a role. Without any restriction, the code will look as follows.

from typing import Dict

roles_str = ["manager", "developer", "tester"]

role_ids_with_str: Dict[str, int] = {
    "manager": 1,
    "developer": 2,
    "tester": 3,
}


def get_role_id_with_str(role: str):
    if role in roles_str:
        print(f"ID ({role}): {role_ids_with_str[role]}")
    else:
        print(f"Role not found: {role}")

get_role_id_with_str("manager") # ID (manager): 1
get_role_id_with_str("owner")   # Role not found: owner

There are two issues here.

First, role_ids_with_str is defined as str which can take any string. It means that we might make a typo in the list and it causes a problem. It can’t be recognized until it’s tested.

Second, the role parameter must be checked if it is one of the acceptable values. It’s not a bad thing to check the value but we want to remove it if possible.

Restrict acceptable values by Enum object

I guess Enum is often used to restrict values in other languages. We can define Enum in Python too. Let’s see the code.

from typing import Dict
from enum import Enum

class Role(Enum):
    Manager = "manager"
    Developer = "developer"
    Tester = "tester"


role_ids_with_enum: Dict[Role, int] = {
    Role.Manager: 1,
    Role.Developer: 2,
    Role.Tester: 3,
}


def get_role_id_with_enum(role: Role):
    print(f"ID ({role.value}): {role_ids_with_enum[role]}")
    print(f"role: ({role}), name: {role.name}, value: {role.value}")

get_role_id_with_enum(Role.Manager) 
# ID (manager): 1
# role: (Role.Manager), name: Manager, value: manager

The two issues are resolved here. To define role_ids_with_enum, we have to use one of the values of Role. IntelliSense will tell us the mistake if we do something wrong. The function requires Enum of Role. Since it’s always Enum, we don’t have to check the value.

However, if we change the data type from string to Enum, the client/caller side must change the code as well. It introduces a breaking change. If using Enum, we should carefully consider potentially affected areas.

Furthermore, if we need the variable name or value, we have to use role.name or role.value. The function above is very small and thus we don’t have to consider the differences in this case but if it is used in the internal functions, there might be several places to change.

Of course, an error is shown if passing a string.

# Argument of type "Literal['developer']" cannot be assigned to parameter "role" of type "Role" in function "get_role_id_with_enum"
# "Literal['developer']" is incompatible with "Role"
get_role_id_with_enum("developer")

Restrict acceptable values by Literal type

Literal is a good choice to avoid a breaking change.

from typing import Dict, Literal

RolesLiteral = Literal["manager", "developer", "tester"]
role_ids_with_literal: Dict[RolesLiteral, int] = {
    "manager": 1,
    "developer": 2,
    "tester": 3,
}


def get_role_id_with_literal(role: RolesLiteral):
    print(f"ID ({role}): {role_ids_with_literal[role]}")

get_role_id_with_literal("manager") # ID (manager): 1

Using Literal resolves all the issues that I explained above. The parameter is still a normal string. Therefore, the client/caller side doesn’t have to change the code.

If the wrong value is added to the dictionary, an error is shown on the line. If we need to expand the list, we can safely add additional items.

# Expression of type "dict[str, int]" cannot be assigned to declared type "Dict[RolesLiteral, int]"
role_ids_with_literal: Dict[RolesLiteral, int] = {
    "manager": 1,
    "developer": 2,
    "tester": 3,
    "owner": 4,
}

It also shows an error if a wrong value is specified to the function.

# Argument of type "Literal['owner']" cannot be assigned to parameter "role" of type "roles_literal" in function "get_role_id_with_literal"
#   Type "Literal['owner']" cannot be assigned to type "roles_literal"
#     "Literal['owner']" cannot be assigned to type "Literal['manager']"
#     "Literal['owner']" cannot be assigned to type "Literal['developer']"
#     "Literal['owner']" cannot be assigned to type "Literal['tester']"
get_role_id_with_literal("owner")

Get all values of the literal type

Use get_args if an object needs to be created dynamically with the values.

from typing import Dict, get_args, Literal

def get_literal_values():
    values = get_args(RolesLiteral)
    print(values)  # ('manager', 'developer', 'tester')
    dynamic_ids = {}
    for index, x in enumerate(values):
        dynamic_ids[x] = index + 1

    print(dynamic_ids) # {'manager': 1, 'developer': 2, 'tester': 3}