Support custom objects Manager

[Javrd]

Hi, Django allow you to define in your class the manager you want to use as default when using objects. This can be useful for adding new methods to your manager, annotating things always in your query, etc...

When using django-types unfortunately it can't infer the type of the custom manager if you have define it in the class. Instead it's always a BaseManager of the model you are using.

This is defined here: https://github.com/sbdchd/django-types/blob/main/django-stubs/db/models/base.pyi#L27C5-L27C41

I've check that deleting that line does correctly infer the type from the custom Manager I'm using, but of course it also lose the Manager type if a class doesn't implement a custom manager objects field.

I've check how do django-stubs solve this and they use some script to transform it with the mypy plugin. I quite new with all this python types and I don't know if it supports somehow to define that the type of the fields objects will be BaseManager OR any class implementing that BaseManager that will be defined later in the class. Not sure if that is even possible because how django does it "magic" for adding the objects field.

It would be great add support to custom objects Manager though, and the only workaround I can use in my project is just use a different manager instead of objects for my custom Managers, but that means a huge refactor.

[arturovergara]

Any update on this?:( Or at least is there any tricky way to define the type for a custom manager ("objects") using pyright? I'd really appreciate any help!

[sbdchd]

Do you have some example code, might be able to figure something out 👀

[patrick91]

@sbdchd here's an example of custom manager: https://testdriven.io/blog/django-custom-user-model/#model-manager

[sbdchd]

hmm would something like this work?

Note the generic passed to the BaseUserManager

class CustomUserManager(BaseUserManager["User"]):
    def create_user(self, email, password, **extra_fields):
        if not email:
            raise ValueError(_("The Email must be set"))
        email = self.normalize_email(email)
        user = self.model(email=email, **extra_fields)
        user.set_password(password)
        user.save()
        return user

    def create_superuser(self, email, password, **extra_fields):
        extra_fields.setdefault("is_staff", True)
        extra_fields.setdefault("is_superuser", True)
        extra_fields.setdefault("is_active", True)

        if extra_fields.get("is_staff") is not True:
            raise ValueError(_("Superuser must have is_staff=True."))
        if extra_fields.get("is_superuser") is not True:
            raise ValueError(_("Superuser must have is_superuser=True."))
        return self.create_user(email, password, **extra_fields)
        
        
class User(models.Model):
    objects = CustomerUserManager()
    # the rest of your fields

[sbdchd]

Another example

class InviteManager(models.Manager["Invite"]):
    def create_invite(
        self, email: str, team: Team, level: str, creator: User
    ) -> Invite:
        user = User.objects.filter(email=email).first()
        if not user:
            user = User.objects.create_user(email=email)
        m = Membership.objects.create(
            user=user, team=team, level=level, is_active=False
        )
        return self.model.objects.create(membership=m, creator=creator)


class Invite(CommonInfo):
    id: int
    membership = models.OneToOneField(Membership, on_delete=models.CASCADE)
    creator = models.ForeignKey(User, on_delete=models.CASCADE)

    OPEN = "open"
    DECLINED = "declined"
    ACCEPTED = "accepted"

    INVITE_STATUS = ((OPEN, OPEN), (DECLINED, DECLINED), (ACCEPTED, ACCEPTED))

    status = models.CharField(max_length=11, choices=INVITE_STATUS, default=OPEN)

    objects = InviteManager()

[Javrd]

The problem in both examples is when you use it. If you hover objects attribute in the model with the mouse in vscode, it will point to the correct custom manager class in each example. However, as soon as you use it in next line
User.objects.create_user(...)
or
Invite.objects.create_invite(...)
the custom function from the custom manager is type "any", without knowledge of function inputs and output. Also, if you hover the objects in that line, it is not pointing to the custom manager class anymore, but to the BaseManager[User] or BaseManager[Invite] in each case.

Note that in the second example you need to be able to resolve the CommonInfo import. If you ignore it and let it undefined, it wont know it comes from models.Model and then it will work well, resolving the custom manager. But as soon as it detect it inherits from models.Model, then when you use it it always changes to the base class because this definition I think.

You can see in the second example that the problem is present where it's used. I've check cloning the whole repo and checking in vscode with lines like this one.

[Javrd]

Related with this probably #194