CODEWITHSUNDEEP
pythondjangodjango-modelsdjango-rest-framework
Jeroen van der Merwe
Jeroen van der Merwe

Reputation: 406

How to add a custom method to a model field in Django?

I have two models that will use the same CardNumberField() to store credit card numbers. How can I add a custom method to the field to mask the card numbers?

I have created the CardNumberField() which inherits from models.Charfield:

# CARD NUMBER FIELD
class CardNumberField(models.CharField):
    description = _('card number')

    def __init__(self, *args, **kwargs):
        kwargs['max_length'] = 19
        super().__init__(*args, **kwargs)

The CardNumberField() is then imported and used in my customers/models.py:

# CARD MODEL
class Card(models.Model):
    number = CardNumberField()
    ...

    def __str__(self):
        return 'Card [{number}]'.format(number=self.number)

...and in my transactions/models.py:

# TRANSACTION MODEL
class Transaction(models.Model):
    card_number = CardNumberField()
    ...

    def __str__(self):
        return 'Transaction ...'

So, how can I add the following method to my CardNumberField() to be used by both of my models?

def masked_number(self):
    # display masked card number
    number = self.number
    return number[-4:].rjust(len(number), '#')

Also, how will I grab this field method in a DRF serializer class?

Upvotes: 2

Views: 1824

Answers (2)

Alain Bianchini
Alain Bianchini

Reputation: 4171

Use an abstract model instead:

class ModelWithCardNumber(models.Model):
   card_number = models.CharField(max_length=19)
   
   @property
   def masked_number(self):
      return self.card_number[-4:].rjust(len(number), '#')
   
   class Meta:
      abstract = True


class Card(ModelWithCardNumber):
    def __str__(self):
        return 'Card [{number}]'.format(number=self.number)


class Transaction(ModelWithCardNumber):
    def __str__(self):
        return 'Transaction ...'

Now in your serializer you can access Card.masked_number and Transaction.masked_number.

Upvotes: 0

willeM_ Van Onsem
willeM_ Van Onsem

Reputation: 476557

You can override the contribute_to_class method to not only contribute the field, but also include an extra method:

from functools import partialmethod

def _mask_number(self, field):
    number = getattr(self, field.attname)
    return number[-4:].rjust(len(number), '#')

# CARD NUMBER FIELD
class CardNumberField(models.CharField):
    description = _('card number')

    def __init__(self, *args, **kwargs):
        kwargs['max_length'] = 19
        super().__init__(*args, **kwargs)

    def contribute_to_class(self, cls, name, **kwargs):
        super().contribute_to_class(cls, name, **kwargs)
        setattr(
            cls, f'masked_{self.name}',
            partialmethod(_mask_number, field=self)
        )

If you add a field foo to a model class, it will automatically add a masked_foo method to that class. This thus also means that if you have two or more CardNumberFields, it will add two or more masked_foo methods.

Upvotes: 3

Related Questions