sms_feed/feed_content/management/commands/import_from_...

from django.core.management.base import BaseCommand, CommandError
from django.db import transaction
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

    with transaction.atomic():
        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)