How can I add python type annotations to the flask global context g?

问题

I have a decorator which adds a user onto the flask global context g:

class User:
    def __init__(self, user_data) -> None:
        self.username: str = user_data["username"]
        self.email: str = user_data["email"]

def login_required(f):
    @wraps(f)
    def wrap(*args, **kwargs):
        user_data = get_user_data()
        user = User(user_data)
        g.user = User(user_data)

        return f(*args, **kwargs)

    return wrap

I want the type (User) of g.user to be known when I access g.user in the controllers. How can I achieve this? (I am using pyright)

回答1:

I had a similar issue described in Typechecking dynamically added attributes. One solution is to add the custom type hints using typing.TYPE_CHECKING:

from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from flask.ctx import _AppCtxGlobals

    class MyGlobals(_AppCtxGlobals):
        user: 'User'

    g = MyGlobals()
else:
    from flask import g

Now e.g.

reveal_type(g.user)

will emit

note: Revealed type is 'myapp.User'

If the custom types should be reused in multiple modules, you can introduce a partial stub for flask. The location of the stubs is dependent on the type checker, e.g. mypy reads custom stubs from the MYPY_PATH environment variable, pyright looks for a typings directory in the project root dir etc. Example of a partial stub:

# _typeshed/flask/__init__.pyi

from typing import Any
from flask.ctx import _AppCtxGlobals
from models import User


def __getattr__(name: str) -> Any: ...  # incomplete


class MyGlobals(_AppCtxGlobals):
    user: User
    def __getattr__(self, name: str) -> Any: ...  # incomplete


g: MyGlobals

回答2:

You could proxy the g object. Consider the following implementation:

import flask


class User:
    ...


class _g:

    user: User
    # Add type hints for other attributes
    # ...

    def __getattr__(self, key):
        return getattr(flask.g, key)


g = _g()

回答3:

You can annotate an attribute on a class, even if that class isn't yours, simply with a colon after it. For example:

g.user: User

That's it. Since it's presumably valid everywhere, I would put it at the top of your code:

from functools import wraps

from flask import Flask, g

app = Flask(__name__)


class User:
    def __init__(self, user_data) -> None:
        self.username: str = user_data["username"]
        self.email: str = user_data["email"]


# Annotate the g.user attribute
g.user: User


def login_required(f):
    @wraps(f)
    def wrap(*args, **kwargs):
        g.user = User({'username': 'wile-e-coyote',
                       'email': 'coyote@localhost'})

        return f(*args, **kwargs)

    return wrap


@app.route('/')
@login_required
def hello_world():
    return f'Hello, {g.user.email}'


if __name__ == '__main__':
    app.run()

That's it.

来源：https://stackoverflow.com/questions/60123011/how-can-i-add-python-type-annotations-to-the-flask-global-context-g

标签

python

Flask

mypy

python-typing

pyright