diff --git a/django_seed/seeder.py b/django_seed/seeder.py index a196871..07d8e97 100644 --- a/django_seed/seeder.py +++ b/django_seed/seeder.py @@ -14,6 +14,17 @@ def __init__(self, model): self.model = model self.field_formatters = {} + @staticmethod + def build_relation(field, related_model): + def func(inserted): + if related_model in inserted and inserted[related_model]: + pk = random.choice(inserted[related_model]) + return related_model.objects.get(pk=pk) + else: + message = 'Field {} cannot be null'.format(field) + raise SeederException(message) + return func + def guess_field_formatters(self, faker): """ Gets the formatter methods for each field using the guessers @@ -28,17 +39,7 @@ def guess_field_formatters(self, faker): for field in model._meta.fields: field_name = field.name if isinstance(field, (ForeignKey, ManyToManyField, OneToOneField)): - related_model = field.rel.to - - def build_relation(inserted): - if related_model in inserted and inserted[related_model]: - pk = random.choice(inserted[related_model]) - return related_model.objects.get(pk=pk) - else: - message = 'Field {} cannot be null'.format(field) - raise SeederException(message) - - formatters[field_name] = build_relation + formatters[field_name] = self.build_relation(field, field.rel.to) continue if isinstance(field, AutoField):