Django migration for TextChoices field

remap_values_migration.py

from django.db.migrations.operations.base import Operation
from django.db import transaction


class Delete:
    pass


"""
RemapValues changes the values of a field
Run this after altering a TextChoices field
Keep in mind that all values will be remapped or defaulted
"""
class RemapValues(Operation):
    reversible = True

    def __init__(self, model_field_pairs, mapping):
        # TODO: the model_field_pairs could be guessed from the TextChoices model that is being altered
        assert model_field_pairs
        self.model_field_pairs = model_field_pairs
        self.mapping = mapping

    def state_forwards(self, app_label, state):
        pass

    def database_forwards(self, app_label, schema_editor, from_state, to_state, mapping=None):
        mapping = mapping or self.mapping
        from_state.clear_delayed_apps_cache()
        db_alias = schema_editor.connection.alias
        for model_name, field_name in self.model_field_pairs:
            model = from_state.apps.get_model(app_label, model_name)
            default = model._meta.get_field(field_name).default
            with transaction.atomic():
                for obj in model.objects.using(db_alias).all():
                    current = getattr(obj, field_name)
                    updated = self.mapping.get(current, default)
                    if updated is Delete:
                        obj.delete()
                    else:
                        setattr(obj, field_name, updated)
                        obj.save()

    def database_backwards(self, app_label, schema_editor, from_state, to_state):
        inverse = {v: k for k, v in self.mapping.items()}
        self.database_forwards(app_label, schema_editor, from_state, to_state, inverse)

    def describe(self):
        return "Remaps " + ", ".join(f"{k} to {v}" for k, v in self.mapping) + " in " + ", ".join(f"{m}.{f}" for m, f in self.model_field_pairs)

    @property
    def migration_name_fragment(self):
        m, f = self.model_field_pairs[0]
        return f"remap_values_{m}_{f}"