Results-Page/apps/home/management/commands/create_sections_fixture.py
2024-11-05 11:48:32 +00:00

79 lines
2.5 KiB
Python

"""Command to create test data fixture for sections."""
import json
from string import ascii_uppercase
from django.core.management.base import BaseCommand
from mainapp.models import Section, SectionValidator
def name_sort_key(name):
if len(name) == 1:
return (0, name)
else:
return (1, name[:-1], name[-1])
class Command(BaseCommand):
help = "Creates a fixture file for Team objects"
def add_arguments(self, parser):
parser.add_argument("amount_of_sections", type=int, help="Number of sections to create")
def handle(self, *args, **options):
amount_of_sections = options["amount_of_sections"]
used_names = set()
last_name = ""
for i in range(amount_of_sections):
add_character = last_name.endswith("Z")
if add_character:
last_name = ""
new_name = last_name if add_character else ""
if not new_name.endswith("Z"):
new_name = new_name[:-1] if len(new_name) > 1 else new_name
try:
new_name += ascii_uppercase[i]
except IndexError:
quotient, remainder = divmod(i, 26)
new_name += ascii_uppercase[quotient - 1] + ascii_uppercase[remainder]
while new_name in used_names:
if new_name[-1] == "Z":
idx = len(new_name) - 2
while idx >= 0 and new_name[idx] == "Z":
idx -= 1
if idx < 0:
new_name = "A" * (len(new_name) + 1)
else:
new_name = new_name[:idx] + chr(ord(new_name[idx]) + 1) + "A" * (len(new_name) - idx - 1)
else:
new_name = new_name[:-1] + chr(ord(new_name[-1]) + 1)
used_names.add(new_name)
last_name = new_name
names = sorted([name for name in used_names], key=name_sort_key)
print([name for name in names])
sections = [Section(name=name) for name in names]
Section.objects.bulk_create(sections)
fixture = [{
"model": "mainapp.section",
"pk": section.pk,
"fields": {
"name": section.name
}
} for section in sections]
for section in sections:
section.delete()
with open("src/mainapp/fixtures/sections_fixture.json", "w") as file:
file.write(json.dumps(fixture, indent=4))
self.stdout.write(self.style.SUCCESS("Created sections_fixture.json."))