How do I properly extend django's built-in User annotations to include relation Managers?

[argilya]

Hello,

The django code base I work with makes use of django's default User model. I would like to annotate this model with relation managers. For example, if there is a table called Posts that has a FK from User, I would like Pyright to recognize that reveal_type(User().posts) is Manger[Post]. I have read your README on how to achieve this with your own models. However, User is defined by django, and so I can't edit it.

How do I annotate djagno's User to let Pyright know that User().posts is legal? This question is probably a bit naive, but I am new to both django and to python types, and so I would appreciate the help!

I came up with 2 ways to solve this problem:

1. Define a superset of User with the manger annotation included, like so:

   class AnnotatedUser(User):
   posts: Manager[Post] # reveal_type() correctly shows this as Manager[Post]

   This approach works, but I am unhappy with it, because I will have to refactor the entire code base to use AnnotatedUser instead of User. I would like to avoid that.

2. So I figured I could just make a custom type stub for User that included posts: Manager[Post]. That's where I ran into problems. I wrote a stub like so:

   
   # /typings/django/contrib/auth/models.pyi
   from __future__ import annotations    

from typing import TYPE_CHECKING

if TYPE_CHECKING:
from posts.models import Post
from django.db.models import Manager from django.contrib.auth.models import AbstractUser

I just copied the User stub from django-stubs and added posts

class User(AbstractUser): objects: Manager[User]
posts: Manager[Post]

So now `User().posts` is properly revealed as `Manger[Posts]`. However, I lose all other types like `password` and `username`, etc, that are defined on `AbstractUser` in `django-stubs` (i.e. `User().password` is now Unknown). Why does this happen? I think I do not fully understand how Pyright uses stubs. Does my little stub overwrite what's present in `django-stubs/.../models.pyi`, and that's why I am not getting any types from AbstractUser?

How should I approach making a custom stub for `User` that adds `Manager[Post]` type but preserves everything else that's already typed in `django-stubs`? Is that possible?

Thank you!

[sbdchd]

Oof this is a tricky problem with the stubs. Basically there isn’t a way to add fields to an existing type with Python’s typing. One option is to create a protocol that mirrors the shape of the User, but then you’ll need to update all the functions that take a User as a param accordingly.

I think your approach of copying all the fields to your own type stub should also work, it’s also pretty manual as you’ve noticed.

[argilya]

Thank you for the response. I'll go with the last option, since I won't have to refactor the code base.