from argparse import FileType from sys import stdin from django.core.management import BaseCommand from media.models import Author, CD class Command(BaseCommand): def add_arguments(self, parser): parser.add_argument('input', nargs='?', type=FileType('r'), default=stdin, help="CD to be imported.") def handle(self, *args, **options): file = options["input"] cds = [] for line in file: cds.append(line[:-1].split('|', 2)) print("Registering", len(cds), "CDs") imported = 0 for cd in cds: if len(cd) != 3: continue title = cd[0] side = cd[1] authors_str = cd[2].split('|') authors = [Author.objects.get_or_create(name=author)[0] for author in authors_str] cd, created = CD.objects.get_or_create( title=title, side_identifier=side, ) cd.authors.set(authors) cd.save() if not created: self.stderr.write(self.style.WARNING( "One CD was already imported. Skipping...")) else: self.stdout.write(self.style.SUCCESS( "CD imported")) imported += 1 self.stdout.write(self.style.SUCCESS( "{count} CDs imported".format(count=imported)))