Как использовать подсказки типа Python с Django QuerySet?

Одним из решений может быть использование класса типизации Union.

from typing import Union, List
from django.db.models import QuerySet
from my_app.models import MyModel

def somefunc(row: Union[QuerySet, List[MyModel]]):
    pass

Теперь, когда вы нарезаете аргумент row, он будет знать, что возвращаемый тип является либо другим списком MyModel, либо экземпляром MyModel, а также намекает, что методы класса QuerySet доступны и для аргумента row.

Существует специальный пакет под названием django-stubs (название следует за PEP561) для ввода django код.

Вот как это работает:

# server/apps/main/views.py
from django.http import HttpRequest, HttpResponse
from django.shortcuts import render

def index(request: HttpRequest) -> HttpResponse:
    reveal_type(request.is_ajax)
    reveal_type(request.user)
    return render(request, 'main/index.html')

Выход:

» PYTHONPATH="$PYTHONPATH:$PWD" mypy server
server/apps/main/views.py:14: note: Revealed type is 'def () -> builtins.bool'
server/apps/main/views.py:15: note: Revealed type is 'django.contrib.auth.models.User'

А с моделями и QuerySetс:

# server/apps/main/logic/repo.py
from django.db.models.query import QuerySet

from server.apps.main.models import BlogPost

def published_posts() -> 'QuerySet[BlogPost]':  # works fine!
    return BlogPost.objects.filter(
        is_published=True,
    )

Выход:

reveal_type(published_posts().first())
# => Union[server.apps.main.models.BlogPost*, None]

• Полное руководство: https://sobolevn.me/2019/08/typechecking-django-and-drf
• Типы для django: https://github.com/typeddjango/django-stubs
• Типы для drf: https://github.com/typeddjango/djangorestframework-stubs

Я сделал этот вспомогательный класс, чтобы получить подсказку общего типа:

from django.db.models import QuerySet
from typing import Iterator, Union, TypeVar, Generic

T = TypeVar("T")

class ModelType(Generic[T]):
    def __iter__(self) -> Iterator[Union[T, QuerySet]]:
        pass

Затем используйте его следующим образом:

def somefunc(row: ModelType[SomeModel]):
    pass

Это уменьшает шум каждый раз, когда я использую этот тип, и делает его пригодным для использования между моделями (например, ModelType[DifferentModel]).

Это улучшенный вспомогательный класс Ор Дуана.

from django.db.models import QuerySet
from typing import Iterator, TypeVar, Generic

_Z = TypeVar("_Z")  

class QueryType(Generic[_Z], QuerySet):
    def __iter__(self) -> Iterator[_Z]: ...

Этот класс используется специально для объекта QuerySet, например, когда вы используете filter в запросе.
Образец:

from some_file import QueryType

sample_query: QueryType[SampleClass] = SampleClass.objects.filter(name=name)

Теперь интерпретатор распознает sample_query как объект QuerySet, и вы получите предложения, такие как count(), и при циклическом просмотре объектов вы получите предложения для SampleClass.

Примечание
Этот формат подсказки типа доступен начиная с python3.6.

Вы также можете использовать django_hint, в котором есть классы подсказок специально для Django.

ИМХО, правильный способ сделать это - определить тип, который наследует QuerySet, и указать общий тип возвращаемого значения для итератора.

    from django.db.models import QuerySet
    from typing import Iterator, TypeVar, Generic, Optional

    T = TypeVar("T")


    class QuerySetType(Generic[T], QuerySet):  # QuerySet + Iterator

        def __iter__(self) -> Iterator[T]:
            pass

        def first(self) -> Optional[T]:
            pass

        # ... add more refinements

Затем вы можете использовать его следующим образом:

users: QuerySetType[User] = User.objects.all()
for user in users:
   print(user.email)  # typing OK!
user = users.first()  # typing OK!

На самом деле вы можете делать то, что хотите, если импортируете модуль аннотаций:

from __future__ import annotations
from django.db import models
from django.db.models.query import QuerySet

class MyModel(models.Model):
    pass

def my_function() -> QuerySet[MyModel]:
    return MyModel.objects.all()

Ни MyPy, ни интерпретатор Python не будут жаловаться или вызывать исключения по этому поводу (проверено на python 3.7). MyPy, вероятно, не сможет проверить его тип, но если все, что вам нужно, это задокументировать тип возвращаемого значения, этого должно быть достаточно.

from typing import Iterable

def func(queryset_or_list: Iterable[MyModel]): 
    pass

И набор запросов, и список экземпляров модели являются итерируемыми объектами.

QuerySet — хороший подход к функциям/методам, возвращающим любой набор запросов любых моделей. Набор запросов Django является итерируемым. Но когда возвращаемый тип очень специфичен для одной модели, может быть лучше использовать QuerySet[Model] вместо QuerySet.

Пример: фильтрация всех активных пользователей компании

import datetime
from django.utils import timezone
from myapp.models import User
from collections.abc import Iterable

def get_active_users(company_id: int) -> QuerySet[User]:
    one_month_ago = (timezone.now() - datetime.timedelta(days=30)).timestamp()
    return User.objects.filter(company_id=company_id, is_active=True, 
                               last_seen__gte=one_month_ago)

Приведенная выше сигнатура функции более читабельна, чем def get_active_users(company_id: int) -> QuerySet:

Iterable[User] также будет работать, проверка типов будет жаловаться, когда возвращаемый набор запросов вызывается другими методами.

def func() -> Iterable[User]:
    return User.objects.all()

users = func()
users.filter(email__startswith='support')

Выход MyPy

"Iterable[User]" has no attribute "filter"

Я обнаружил, что использую typing.Sequence для решения аналогичной проблемы:

from typing import Sequence


def print_emails(users: Sequence[User]):
    for user in users:
        print(user.email)


users = User.objects.all()


print_emails(users=users)

Насколько я знаю из документов:

Последовательность — это все, что поддерживает len() и .getitem(), независимо от ее фактического типа.