from django.core.management.base import BaseCommand, CommandError import xlrd from ... import models from pick import pick class SheetWrapper(object): def __init__(self, sheet): self.sheet = sheet def __str__(self): return self.sheet.name def import_sheet(sheet, category): start_row = None start_col = None for row in range(sheet.nrows): for col in range(sheet.ncols): cell = sheet.cell(row, col) if cell.ctype == xlrd.XL_CELL_TEXT: if cell.value.strip().lower() == 'day': start_row = row + 1 start_col = col break if start_row is None: raise RuntimeError("No 'Day' column found in sheet %s" % sheet.name) if start_row >= sheet.nrows: raise RuntimeError("'Day' column in sheet %s contains no data" % sheet.name) if start_col >= sheet.ncols: raise RuntimeError("No columns found to the right of the 'Day' column in sheet %s" % sheet.name) expected_day = 1 end_row = None feed_items = [] for row in range(start_row, sheet.nrows): cell = sheet.cell(row, start_col) if cell.ctype != xlrd.XL_CELL_NUMBER: break day = int(cell.value) if day != expected_day: raise RuntimeError("'Day' column in sheet %s contains an unexpected value: %s (expecting %s)" % (sheet.name, day, expected_day)) cell = sheet.cell(row, start_col + 1) if cell.ctype != xlrd.XL_CELL_TEXT: raise RuntimeError("No text found for day %s in sheet %s" % (day, sheet.name)) feed_items.append({'day': day, 'message_text': cell.value}) expected_day += 1 end_row = row if end_row + 1 < sheet.nrows: cell = sheet.cell(end_row + 1, start_col) if cell.ctype != xlrd.XL_CELL_EMPTY: raise RuntimeError("'Day' column in sheet %s contains a non-numerical value: %s" % (sheet.name, cell.value)) option, index = pick(['No', 'Yes'], title='You are about to import days %s to %s\nfrom sheet %s\ninto category %s.\nContinue?' % (feed_items[0]['day'], feed_items[-1]['day'], sheet.name, category.name)) if index == 0: return models.FeedItem.objects.filter(feed_category=category).delete() for feed_item in feed_items: models.FeedItem.objects.create( feed_category=category, day_number=feed_item['day'], message_text=feed_item['message_text'] ) class Command(BaseCommand): help = 'Imports feed items from an Excel spreedsheet.' def add_arguments(self, parser): parser.add_argument('file_name') def handle(self, *args, **options): file_name = options['file_name'] category_choices = ['>> Quit'] for category in models.FeedCategory.objects.order_by('name'): category_choices.append(category) book = xlrd.open_workbook(file_name) sheet_choices = ['>> Quit'] for sheet in book.sheets(): sheet_choices.append(SheetWrapper(sheet)) while True: option, index = pick(sheet_choices, title='Select the worksheet to be imported:') if index == 0: break sheet = sheet_choices[index].sheet option, index = pick(category_choices, title='Select the destination category:') if index == 0: break category = category_choices[index] import_sheet(sheet, category)