feat: add support for integer flags

Author: bellini666

.coveragerc

@@ -4,3 +4,11 @@ omit = .venv/**
 
 [report]
 precision = 2
+exclude_lines =
+    pragma: nocover
+    pragma:nocover
+    if TYPE_CHECKING:
+    @overload
+    @abstractmethod
+    @abc.abstractmethod
+    assert_never

README.md

@@ -17,8 +17,10 @@ pip install django-choices-field
 ## Usage
 
 ```python
+import enum
+
 from django.db import models
-from django_choices_field import TextChoicesField, IntegerChoicesField
+from django_choices_field import TextChoicesField, IntegerChoicesField, IntegerChoicesFlag
 
 
 class MyModel(models.Model):
@@ -30,23 +32,41 @@ class MyModel(models.Model):
         FIRST = 1, "First Description"
         SECOND = 2, "Second Description"
 
-    c_field = TextChoicesField(
+    class IntegerFlagEnum(IntegerChoicesFlag):
+        FIRST = enum.auto(), "First Option"
+        SECOND = enum.auto(), "Second Option"
+        THIRD = enum.auto(), "Third Option"
+
+    text_field = TextChoicesField(
         choices_enum=TextEnum,
         default=TextEnum.FOO,
     )
-    i_field = IntegerChoicesField(
+    integer_field = IntegerChoicesField(
         choices_enum=IntegerEnum,
         default=IntegerEnum.FIRST,
     )
+    flag_field = IntegerChoicesFlagField(
+        choices_enum=IntegerFlagEnum,
+        default=IntegerFlagEnum.FIRST | IntegerFlagEnum.SECOND,
+    )
 
 
 obj = MyModel()
-obj.c_field  # MyModel.TextEnum.FOO
-isinstance(obj.c_field, MyModel.TextEnum) # True
-obj.i_field  # MyModel.IntegerEnum.FIRST
-isinstance(obj.i_field, MyModel.IntegerEnum) # True
+reveal_type(obj.text_field)  # MyModel.TextEnum.FOO
+assert isinstance(obj.text_field, MyModel.TextEnum)
+assert obj.text_field == "foo"
+
+reveal_type(obj.integer_field)  # MyModel.IntegerEnum.FIRST
+assert isinstance(obj.integer_field, MyModel.IntegerEnum)
+assert obj.integer_field == 1
+
+reveal_type(obj.flag_field)  # MyModel.IntegerFlagEnum.FIRST | MyModel.IntegerFlagEnum.SECOND
+assert isinstance(obj.integer_field, MyModel.IntegerFlagEnum)
+assert obj.flag_field == 3
 ```
 
+NOTE: The `IntegerChoicesFlag` requires python 3.11+ to work properly.
+
 ## License
 
 This project is licensed under MIT licence (see `LICENSE` for more info)

django_choices_field/__init__.py

@@ -1,6 +1,8 @@
 from .fields import IntegerChoicesField, TextChoicesField
+from .types import IntegerChoicesFlag
 
 __all__ = [
+    "IntegerChoicesFlag",
     "TextChoicesField",
     "IntegerChoicesField",
 ]

django_choices_field/fields.py

@@ -1,8 +1,12 @@
+import functools
+import itertools
 from typing import ClassVar, Dict, Optional, Type
 
 from django.core.exceptions import ValidationError
 from django.db import models
 
+from .types import IntegerChoicesFlag
+
 
 class TextChoicesField(models.CharField):
     description: ClassVar[str] = "TextChoices"
@@ -97,3 +101,65 @@ def formfield(self, **kwargs):  # pragma:nocover
                 **kwargs,
             },
         )
+
+
+class IntegerChoicesFlagField(models.IntegerField):
+    description: ClassVar[str] = "IntegerChoicesFlag"
+    default_error_messages: ClassVar[Dict[str, str]] = {
+        "invalid": "“%(value)s” must be a subclass of %(enum)s.",
+    }
+
+    def __init__(
+        self,
+        choices_enum: Type[IntegerChoicesFlag],
+        verbose_name: Optional[str] = None,
+        name: Optional[str] = None,
+        **kwargs,
+    ):
+        self.choices_enum = choices_enum
+
+        default_choices = choices_enum.choices
+        kwargs["choices"] = default_choices[:]
+        for i in range(1, len(default_choices)):
+            for combination in itertools.combinations(default_choices, i + 1):
+                kwargs["choices"].append(
+                    (
+                        functools.reduce(lambda a, b: a | b[0], combination, 0),
+                        "|".join(c[1] for c in combination),
+                    ),
+                )
+
+        super().__init__(verbose_name=verbose_name, name=name, **kwargs)
+
+    def deconstruct(self):
+        name, path, args, kwargs = super().deconstruct()
+        kwargs["choices_enum"] = self.choices_enum
+        return name, path, args, kwargs
+
+    def to_python(self, value):
+        if value is None:
+            return None
+
+        try:
+            return self.choices_enum(int(value) if isinstance(value, str) else value)
+        except ValueError as e:
+            raise ValidationError(
+                self.error_messages["invalid"],
+                code="invalid",
+                params={"value": value, "enum": self.choices_enum},
+            ) from e
+
+    def from_db_value(self, value, expression, connection):
+        return self.to_python(value)
+
+    def get_prep_value(self, value):
+        value = super().get_prep_value(value)
+        return self.to_python(value)
+
+    def formfield(self, **kwargs):  # pragma:nocover
+        return super().formfield(
+            **{
+                "coerce": self.to_python,
+                **kwargs,
+            },
+        )

django_choices_field/fields.pyi

@@ -13,6 +13,8 @@ from typing import (
 from django.db.models import Field, IntegerChoices, TextChoices
 from typing_extensions import TypeAlias
 
+from django_choices_field.types import IntegerChoicesFlag
+
 _ValidatorCallable: TypeAlias = Callable[..., None]
 _ErrorMessagesToOverride: TypeAlias = Dict[str, Any]
 
@@ -145,3 +147,68 @@ class IntegerChoicesField(Generic[_I], Field[_I, _I]):
         allow_files: bool = ...,
         allow_folders: bool = ...,
     ) -> IntegerChoicesField[_I | None]: ...
+
+_IF = TypeVar("_IF", bound=Optional[IntegerChoicesFlag])
+
+class IntegerChoicesFlagField(Generic[_IF], Field[_IF, _IF]):
+    choices_enum: type[_IF]
+    @overload
+    def __new__(
+        cls,
+        choices_enum: type[_IF],
+        verbose_name: str | None = ...,
+        name: str | None = ...,
+        primary_key: bool = ...,
+        max_length: int | None = ...,
+        unique: bool = ...,
+        blank: bool = ...,
+        null: Literal[False] = ...,
+        db_index: bool = ...,
+        default: _IF | Callable[[], _IF] = ...,
+        editable: bool = ...,
+        auto_created: bool = ...,
+        serialize: bool = ...,
+        unique_for_date: str | None = ...,
+        unique_for_month: str | None = ...,
+        unique_for_year: str | None = ...,
+        help_text: str = ...,
+        db_column: str | None = ...,
+        db_tablespace: str | None = ...,
+        validators: Iterable[_ValidatorCallable] = ...,
+        error_messages: _ErrorMessagesToOverride | None = ...,
+        path: str | Callable[..., str] = ...,
+        match: str | None = ...,
+        recursive: bool = ...,
+        allow_files: bool = ...,
+        allow_folders: bool = ...,
+    ) -> IntegerChoicesFlagField[_IF]: ...
+    @overload
+    def __new__(
+        cls,
+        choices_enum: type[_IF],
+        verbose_name: str | None = ...,
+        name: str | None = ...,
+        primary_key: bool = ...,
+        max_length: int | None = ...,
+        unique: bool = ...,
+        blank: bool = ...,
+        null: Literal[True] = ...,
+        db_index: bool = ...,
+        default: _IF | Callable[[], _IF] | None = ...,
+        editable: bool = ...,
+        auto_created: bool = ...,
+        serialize: bool = ...,
+        unique_for_date: str | None = ...,
+        unique_for_month: str | None = ...,
+        unique_for_year: str | None = ...,
+        help_text: str = ...,
+        db_column: str | None = ...,
+        db_tablespace: str | None = ...,
+        validators: Iterable[_ValidatorCallable] = ...,
+        error_messages: _ErrorMessagesToOverride | None = ...,
+        path: str | Callable[..., str] = ...,
+        match: str | None = ...,
+        recursive: bool = ...,
+        allow_files: bool = ...,
+        allow_folders: bool = ...,
+    ) -> IntegerChoicesFlagField[_IF | None]: ...

django_choices_field/types.py

@@ -0,0 +1,29 @@
+import enum
+import sys
+from typing import TYPE_CHECKING
+
+from django.db import models
+from typing_extensions import Self
+
+
+class IntegerChoicesFlag(models.IntegerChoices, enum.Flag):
+    """Enumerated integer choices."""
+
+    if TYPE_CHECKING:
+
+        def __or__(self, other: Self) -> Self:
+            ...
+
+        def __and__(self, other: Self) -> Self:
+            ...
+
+        def __xor__(self, other: Self) -> Self:
+            ...
+
+        def __invert__(self) -> Self:
+            ...
+
+        if sys.version_info >= (3, 11):
+            __ror__ = __or__
+            __rand__ = __and__
+            __rxor__ = __xor__

poetry.lock

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "django-choices-field"
-version = "2.1.2"
+version = "2.2.0"
 description = "Django field that set/get django's new TextChoices/IntegerChoices enum."
 authors = ["Thiago Bellini Ribeiro <[email protected]>"]
 license = "MIT"
@@ -32,6 +32,7 @@ packages = [{ include = "django_choices_field" }]
 [tool.poetry.dependencies]
 python = "^3.8"
 django = ">=3.2"
+typing_extensions = ">=4.0.0"
 
 [tool.poetry.dev-dependencies]
 black = "^23.3.0"

pyproject.toml

@@ -1,6 +1,11 @@
+import enum
+import sys
+
 from django.db import models
 
 from django_choices_field import IntegerChoicesField, TextChoicesField
+from django_choices_field.fields import IntegerChoicesFlagField
+from django_choices_field.types import IntegerChoicesFlag
 
 
 class MyModel(models.Model):
@@ -12,6 +17,11 @@ class IntegerEnum(models.IntegerChoices):
         I_FOO = 1, "I Foo Description"
         I_BAR = 2, "I Bar Description"
 
+    class IntegerFlagEnum(IntegerChoicesFlag):
+        IF_FOO = enum.auto() if sys.version_info >= (3, 11) else 1, "IF Foo Description"
+        IF_BAR = enum.auto() if sys.version_info >= (3, 11) else 2, "IF Bar Description"
+        IF_BIN = enum.auto() if sys.version_info >= (3, 11) else 4, "IF Bin Description"
+
     objects = models.Manager["MyModel"]()
 
     c_field = TextChoicesField(
@@ -30,3 +40,11 @@ class IntegerEnum(models.IntegerChoices):
         choices_enum=IntegerEnum,
         null=True,
     )
+    if_field = IntegerChoicesFlagField(
+        choices_enum=IntegerFlagEnum,
+        default=IntegerFlagEnum.IF_FOO,
+    )
+    if_field_nullable = IntegerChoicesFlagField(
+        choices_enum=IntegerFlagEnum,
+        null=True,
+    )