use django-stubs during static type checking, fixes #60

Author: bckohan

.gitignore

@@ -133,3 +133,4 @@ dmypy.json
 test.db
 django_enum/tests/edit_tests/migrations/00*.py
 benchmark.db
+type_check.py

django_enum/choices.py

@@ -9,21 +9,19 @@
 from django.db.models import Choices
 from django.db.models import IntegerChoices as DjangoIntegerChoices
 from django.db.models import TextChoices as DjangoTextChoices
-
-try:
-    from django.db.models.enums import ChoicesType
-except ImportError:  # pragma: no cover
-    from django.db.models.enums import ChoicesMeta as ChoicesType
+from django.db.models import enums as model_enums
 
 from django_enum.utils import choices, names
 
+ChoicesType = getattr(model_enums, "ChoicesType", getattr(model_enums, "ChoicesMeta"))
+
 DEFAULT_BOUNDARY = getattr(enum, "KEEP", None)
 
 
 try:
     from enum_properties import DecomposeMixin, EnumPropertiesMeta, SymmetricMixin
 
-    class DjangoEnumPropertiesMeta(EnumPropertiesMeta, ChoicesType):
+    class DjangoEnumPropertiesMeta(EnumPropertiesMeta, ChoicesType):  # type: ignore
         """
         A composite meta class that combines Django's Choices metaclass with
         enum-properties metaclass. This metaclass will add Django's expected
@@ -37,15 +35,15 @@ def names(cls):
             For some eccentric enums list(Enum) is empty, so we override names
             if empty
             """
-            return ChoicesType.names.fget(cls) or names(cls, override=True)
+            return super().names or names(cls, override=True)
 
         @property
         def choices(cls):
             """
             For some eccentric enums list(Enum) is empty, so we override
             choices if empty
             """
-            return ChoicesType.choices.fget(cls) or choices(cls, override=True)
+            return super().choices or choices(cls, override=True)
 
     class DjangoSymmetricMixin(SymmetricMixin):
         """

django_enum/fields.py

@@ -31,22 +31,16 @@
     Q,
     SmallIntegerField,
     TimeField,
+    expressions,
 )
 from django.db.models.constraints import CheckConstraint
+from django.db.models.fields import BLANK_CHOICE_DASH
 from django.db.models.query_utils import DeferredAttribute
 from django.utils.deconstruct import deconstructible
 from django.utils.duration import duration_string
 from django.utils.functional import cached_property
 from django.utils.translation import gettext_lazy as _
 
-try:
-    from django.db.models.expressions import DatabaseDefault
-except ImportError:  # pragma: no cover
-
-    class DatabaseDefault:  # type: ignore
-        """Spoof DatabaseDefault for Django < 5.0"""
-
-
 from django_enum.forms import (
     EnumChoiceField,
     EnumFlagField,
@@ -67,14 +61,21 @@ class DatabaseDefault:  # type: ignore
     with_typehint,
 )
 
+
+class _DatabaseDefault:
+    """Spoof DatabaseDefault for Django < 5.0"""
+
+
+DatabaseDefault = getattr(expressions, "DatabaseDefault", _DatabaseDefault)
+
 CONFORM: Union[Enum, Type[NOT_PROVIDED]]
 EJECT: Union[Enum, Type[NOT_PROVIDED]]
 STRICT: Union[Enum, Type[NOT_PROVIDED]]
 
-if sys.version_info >= (3, 11):  # pragma: no cover
+if sys.version_info >= (3, 11):
     from enum import CONFORM, EJECT, STRICT
 else:
-    CONFORM = EJECT = STRICT = NOT_PROVIDED  # pragma: no cover
+    CONFORM = EJECT = STRICT = NOT_PROVIDED
 
 
 MAX_CONSTRAINT_NAME_LENGTH = 64
@@ -624,7 +625,7 @@ def get_default(self) -> Any:
             return 0
         return super().get_default()
 
-    def validate(self, value: Any, model_instance: Model):
+    def validate(self, value: Any, model_instance: Optional[Model]):
         """
         Validates the field as part of model clean routines. Runs super class
         validation routines then tries to convert the value to a valid
@@ -683,14 +684,23 @@ def formfield(self, form_class=None, choices_form_class=None, **kwargs):
         form_field.primitive = self.primitive
         return form_field
 
-    def get_choices(self, **kwargs):  # pylint: disable=W0221
+    def get_choices(
+        self,
+        include_blank=True,
+        blank_choice=tuple(BLANK_CHOICE_DASH),
+        limit_choices_to=None,
+        ordering=(),
+    ):  # pylint: disable=W0221
         if self.enum and issubclass(self.enum, Flag):
-            kwargs["blank_choice"] = [
-                (self.enum(0), "---------")  # pylint: disable=E1102
-            ]
+            blank_choice = [(self.enum(0), "---------")]  # pylint: disable=E1102
         return [
             (getattr(choice, "value", choice), label)
-            for choice, label in super().get_choices(**kwargs)
+            for choice, label in super().get_choices(
+                include_blank=include_blank,
+                blank_choice=list(blank_choice),
+                limit_choices_to=limit_choices_to,
+                ordering=ordering,
+            )
         ]
 
     @staticmethod
@@ -716,15 +726,17 @@ def constraint_name(
             return name[len(name) - MAX_CONSTRAINT_NAME_LENGTH :]
         return name
 
-    def contribute_to_class(self, cls, name, **kwargs):  # pylint: disable=W0221
-        super().contribute_to_class(cls, name, **kwargs)
+    def contribute_to_class(
+        self, cls: Type[Model], name: str, private_only: bool = False
+    ):  # pylint: disable=W0221
+        super().contribute_to_class(cls, name, private_only=private_only)
         if self.constrained and self.enum and issubclass(self.enum, IntFlag):
             # It's possible to declare an IntFlag field with negative values -
             # these enums do not behave has expected and flag-like DB
             # operations are not supported, so they are treated as normal
             # IntEnum fields, but the check constraints are flag-like range
             # constraints, so we bring those in here
-            FlagField.contribute_to_class(self, cls, name, call_base=False, **kwargs)
+            FlagField.contribute_to_class(self, cls, name, private_only=private_only)
         elif self.constrained and self.enum:
             constraint = Q(
                 **{
@@ -1113,7 +1125,7 @@ class FlagField(with_typehint(IntEnumField)):  # type: ignore
     enum: Type[Flag]
 
     def contribute_to_class(
-        self, cls: Type[EnumField], name: str, call_base: bool = True, **kwargs
+        self, cls: Type[Model], name: str, private_only: bool = False
     ) -> None:
         """
         Add check constraints that honor flag fields range and boundary
@@ -1175,8 +1187,10 @@ def contribute_to_class(
                 cls._meta.original_attrs.setdefault(  # pylint: disable=W0212
                     "constraints", cls._meta.constraints  # pylint: disable=W0212
                 )
-        if call_base:
-            IntegerField.contribute_to_class(self, cls, name, **kwargs)
+        if isinstance(self, FlagField):
+            # this may have been called by a normal EnumField to bring in flag-like constraints
+            # for non flag fields
+            IntegerField.contribute_to_class(self, cls, name, private_only=private_only)
 
 
 class SmallIntegerFlagField(FlagField, EnumPositiveSmallIntegerField):
@@ -1275,17 +1289,17 @@ def from_db_value(
             connection,
         )
 
-    def contribute_to_class(self, cls, name, **kwargs):
-        BinaryField.contribute_to_class(self, cls, name, **kwargs)
+    def contribute_to_class(self, cls, name, private_only: bool = False):
+        BinaryField.contribute_to_class(self, cls, name, private_only=private_only)
 
 
 class ExtraBigIntegerFlagField(FlagField, EnumExtraBigIntegerField):
     """
     Flag fields that require more than 64 bits.
     """
 
-    def contribute_to_class(self, cls, name, call_base=True, **kwargs):
-        BinaryField.contribute_to_class(self, cls, name, **kwargs)
+    def contribute_to_class(self, cls, name, private_only: bool = False):
+        BinaryField.contribute_to_class(self, cls, name, private_only=private_only)
 
 
 # ExtraBigIntegerFlagField.register_lookup(HasAnyFlagsExtraBigLookup)

django_enum/filters.py

@@ -3,7 +3,6 @@
 from typing import Tuple, Type
 
 from django.db.models import Field as ModelField
-from django.forms.fields import Field as FormField
 
 from django_enum.forms import EnumChoiceField
 from django_enum.utils import choices
@@ -40,7 +39,7 @@ class Color(TextChoices, s('rgb'), s('hex', case_fold=True)):
         :param kwargs: Any additional arguments for base classes
         """
 
-        field_class: FormField = EnumChoiceField
+        field_class = EnumChoiceField
 
         def __init__(self, *, enum, strict=False, **kwargs):
             self.enum = enum