Idok
Idok

Reputation: 4152

Type hint for dictionary with known keys

I'm using Python 3 typing feature for better autocomplete.

Many times I have functions that return key/value (dictionary) with specific keys. super simple example:

def get_info(name):
    name_first_letter = name[0]
    return {'my_name': name, 'first_letter': name_first_letter}

I want to add type hinting to this function to tell others who use this function what to expect.

I can do something like:

class NameInfo(object):
    def __init__(self, name, first_letter):
        self.name = name
        self.first_letter = first_letter

and then change the function signature to:

def get_info(name) -> NameInfo:

But it requires too much code for each dictionary.

What is the best practice in that case?

Upvotes: 115

Views: 79340

Answers (3)

matt91t
matt91t

Reputation: 181

Use Literal

I think, better way is to use simple from typing import Literal:

def get_info(name) -> dict[Literal["my_name", "first_letter"], str]:
    name_first_letter = name[0]
    return {'my_name': name, 'first_letter': name_first_letter}

edit: thanks to @mkrieger1, Literal["my_name", "first_letter"] means a subset of the {"my_name", "first_letter"}.

Upvotes: 1

cottontail
cottontail

Reputation: 23361

TypedDict can also be created using a typename and a dictionary where the keys are the keys and the values are the type of the values, respectively, of the dict to be returned by get_info(). This allows you to create keys that uses characters not in A-Za-z0-9_.

from typing import TypedDict

NameInfo = TypedDict('NameInfo', {'name': str, 'first-letter': str})

def get_info(name: str) -> NameInfo:
    return {'name': name, 'first-letter': name[0]}

x = get_info('cottontail')
print(x['first-letter'])   # c

Then on PyCharm, the key hint is there (also works on jupyter notebook too by pressing tab).

pycharm

Upvotes: 25

Bernhard
Bernhard

Reputation: 2251

As pointed out by Blckknght, you and Stanislav Ivanov in the comments, you can use NamedTuple:

from typing import NamedTuple


class NameInfo(NamedTuple):
    name: str
    first_letter: str


def get_info(name: str) -> NameInfo:
    return NameInfo(name=name, first_letter=name[0])

Starting from Python 3.8 you can use TypedDict which is more similar to what you want:

from typing import TypedDict


class NameInfo(TypedDict):
    name: str
    first_letter: str


def get_info(name: str) -> NameInfo:
    return {'name': name, 'first_letter': name[0]}

Upvotes: 130

Related Questions