Python: type annotate a function return value that returns a class defined in that function

348 views Asked by At

I'm wondering what the proper way to type annotate a return value of a function or method where a class is defined inside of that function.

For example:

from typing import List
from dataclasses import dataclass

def get_users() -> List['ReturnUser']:

    @dataclass
    class ReturnUser:
       first_name: str
       last_name: str

    return [
        ReturnUser("John", "Doe"), 
        ReturnUser("Jane", "Doe")]

... the ReturnUser dataclass is only relevant within the scope of the function. (I don't care about the object outside of the get_users() function other than that I can access its properties in the calling code).

However, Pylance shows the return value annotation List['ReturnUser'] as a type error. My environment is unfortunately Python 3.7, but also curious about if newer approaches exist in newer Python versions?

How do I annotate a return value that doesn't yet exist like this?

1

There are 1 answers

3
Silvio Mayolo On BEST ANSWER

Not only does ReturnUser not exist yet, but it's a different class every time. We can verify this easily.

>>> a = get_users()
>>> b = get_users()
>>> type(a[0]) is type(b[0])
False

So my first advice is to not do this. Make a module-level class and prefix it with an underscore to indicate that it's private. There are very, very, very few good reasons to define a class inside of a def in Python.

That being said, let's suppose this truly is the only way. Let's suppose there's some design constraint not being shown here, and you can't get rid of that class pattern. As I said, ReturnUser is not a correct return type, since the function returns a value of a different (unrelated) class every time. If you want the return type of your function to be "something with a first_name and last_name field", we can use Protocol to do that.

class _ReturnUserLike(Protocol):
    first_name: str
    last_name: str

def get_users() -> List[_ReturnUserLike]:
    # ... Same class body, including the local class ...

Now the only type-level guarantee to your callers is that you have properties first_name and last_name which are strings. A Protocol is a structural type, which means anything that looks like that class counts as an instance for the purposes of type-checking.