Skip to content

Instantly share code, notes, and snippets.

@douglasmiranda
Created January 4, 2020 17:00
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save douglasmiranda/ac4a8066a9b28b1255bcafab57549e28 to your computer and use it in GitHub Desktop.
Save douglasmiranda/ac4a8066a9b28b1255bcafab57549e28 to your computer and use it in GitHub Desktop.
Django Field.choices alternative

Django Field.choices

I was never very fond of using that tuple of tuples and acessing those choices with something like choices[0].

In my latest projects I was using a solution of my own.

This:

class Choices:
    """An alternative to the classic Choices

    Usage

        class Article(models.Model):
            STATUSES = Choices(DRAFT=_("Draft"), PUBLIC=_("Public"))

            title = models.CharField(_("Title"), max_length=140, db_index=True)
            status = models.CharField(
                max_length=10, choices=STATUSES, default=STATUSES.DRAFT, db_index=True
            )

        Comparison example:

        article = Article.objects.get(pk=1)
        # instance access:
        if article.status == article.STATUSES.PUBLIC:
            publish()

        # class access:
        Article.STATUSES.DRAFT == "DRAFT"

    - As you can see you can just do `choices=STATUSES`, that's because
        `Choices` has iterator capabilities.
    - I did some operator overloading too, so you can directly compare if
        `key == STATUSES.<STATUS>`.

    NOTE: This doesn't work if you want named groups.
    """

    def __init__(self, **kwargs):
        self.choices = kwargs
        self._iter = self.iterator()

    def iterator(self):
        return (choice for choice in self.choices.items())

    def __next__(self):
        return next(self._iter)

    def __iter__(self):
        return self.iterator()

    def __contains__(self, key):
        return key in self.choices

    def __getattr__(self, key):
        if key in self.choices:
            return key
        raise AttributeError(key)


## TESTING:


def test_Choices_initialization():
    STATUSES = Choices(DRAFT="Draft", PUBLIC="Public")


def test_Choices_should_use_keys_when_comparing():
    STATUSES = Choices(DRAFT="Draft", PUBLIC="Public")
    assert STATUSES.DRAFT == "DRAFT"
    assert STATUSES.PUBLIC == "PUBLIC"


def test_Choices_should_use_keys_when_assigning():
    STATUSES = Choices(DRAFT="Draft", PUBLIC="Public")
    draft = STATUSES.DRAFT
    assert STATUSES.DRAFT == draft


def test_Choices_should_expose_dict():
    STATUSES = Choices(DRAFT="Draft", PUBLIC="Public")
    assert STATUSES.choices == {"DRAFT": "Draft", "PUBLIC": "Public"}


def test_Choices_should_be_converted_to_classic_choices():
    STATUSES = Choices(DRAFT="Draft", PUBLIC="Public")
    assert tuple(STATUSES) == (("DRAFT", "Draft"), ("PUBLIC", "Public"))


# Now testing some extra goodies:


def test_Choices_should_have_membership_capabilities():
    STATUSES = Choices(DRAFT="Draft", PUBLIC="Public")
    assert "DRAFT" in STATUSES
    assert "PUBLIC" in STATUSES


def test_Choices_should_have_basic_iterator_capabilities():
    STATUSES = Choices(DRAFT="Draft", PUBLIC="Public")
    assert next(STATUSES) == ("DRAFT", "Draft")
    assert next(STATUSES) == ("PUBLIC", "Public")

    with pytest.raises(StopIteration):
        # Only two elements PUBLIC and DRAFT, so the third time, should be an StopIteration
        # This makes you get out of loops
        assert next(STATUSES)

    iter_again = STATUSES.iterator()
    for item in iter_again:
        pass


def test_Choices_should_unpack_nicely():
    STATUSES = Choices(DRAFT="Draft", PUBLIC="Public")

    key, value = next(STATUSES)
    assert key, value == ("DRAFT", "Draft")
    assert key, value == ("PUBLIC", "Public")

Fortunatelly Django now (3.0) has a new way of dealing with choices. Enumeration Types.

from django.utils.translation import gettext_lazy as _

class Student(models.Model):

    class YearInSchool(models.TextChoices):
        FRESHMAN = 'FR', _('Freshman')
        SOPHOMORE = 'SO', _('Sophomore')
        JUNIOR = 'JR', _('Junior')
        SENIOR = 'SR', _('Senior')
        GRADUATE = 'GR', _('Graduate')

    year_in_school = models.CharField(
        max_length=2,
        choices=YearInSchool.choices,
        default=YearInSchool.FRESHMAN,
    )

    def is_upperclass(self):
        return self.year_in_school in {YearInSchool.JUNIOR, YearInSchool.SENIOR}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment