Using Python Enum and IntEnum Like An Expert

by

Hello guys. If you use forms on your Django web applications and you want to clean the model scheme, you should use enums!

Let’s start

For example, you have a product model and your visitors select something on a form like status.

We consider the options as follows:

  • Task
  • Pending
  • Published

…you have two options. The first one is complex and unmanageable. You can create a tuple like the below and add ‘choices’ attribution to CharField:

STATUS = (
  ("TASK", "Task"),
  ("PENDING", "Pending"),
  ("PUBLISHED", "Published"),
)

or

STATUS = (
  (1, "TASK"),
  (2, "PENDING"),
  (3, "PUBLISHED"),
)

Your product model:

class Product(models.Model):
  category = models.ForeignKey("Category", on_delete=models.CASCADE,
    verbose_name="Category Name", blank=True, null=True)
  title = models.CharField(max_length=50, verbose_name="Product Title",
    blank=True, null=True)
  status = models.CharField(choices=STATUS)

It is correct and will work without errors. Python will save them to db as a string. But it is both complex and unmanageable. Because If you want to use this field anywhere, you need to call like below because a tuple is not useful for calling:

def example(self):
  if self.status == 1:
    make_something()
  elif self.status == 2:
    make_something()

I prefer using Enum or IntEnum. First we need to import Enum:

from enum import Enum, IntEnum

Let’s create our enum classes. Don’t worry, we don’t need to create new migrations.

class Status(enum.Enum):
  TASK = "Task"
  PENDING = "Pending"
  PUBLISHED = "Published"

or (how you want to keep them)

class Status(enum.IntEnum):
  TASK = 1
  PENDING = 2
  PUBLISHED = 3

And your model should be below:

class Product(models.Model):
  ...
  status = models.CharField(
    choices=(
        (x.value, x.name.title()) for x in Status),
    default=Status.TASK
    )

or

class Product(models.Model):
  ...
  status = models.IntegerField(
    choices=(
        (x.value, x.name.title()) for x in Status),
    default=Status.TASK
    )

You can set a default value as I added above. How do you use Enum on functions or anywhere?

You will call enums like below:

def example(self):
  if self.status == Status.TASK.value:
    make_something()
  elif self.status == Status.PUBLISHED.value:
    make_something()

or you can iterate them like below:

def get_states(self):
  return [{"id": s.id, "name": s.name} for s in Status]

Congratz! It’s the Pythonic way.

Leave a Comment