shift_export.py

export_shifts.py

import datetime
import operator
import functools
import os

from unicodecsv import UnicodeWriter

from django.contrib.auth import get_user_model

User = get_user_model()

from shifts.models import Shift

from departments.models import Department

def format_datetime(shift, source, format="%Y-%m-%d %H:%m"):
    when = getattr(shift, source)
    return when.strftime(format)

HEADERS = (
    ('shift_id', operator.attrgetter('pk')),
    ('start_time', functools.partial(format_datetime, source='start_time')),
    ('end_time', functools.partial(format_datetime, source='end_time')),
    ('duration', operator.attrgetter('shift_length')),
    ('department_id', operator.attrgetter('department.pk')),
    ('department_name', operator.attrgetter('department.name')),
    ('full_name', operator.attrgetter('owner.profile.full_name')),
    ('display_name', operator.attrgetter('owner.profile.display_name')),
    ('email', operator.attrgetter('owner.email')),
)


ORDER_BY = ('department', 'owner', 'start_time', 'shift_length')


def shift_to_row(shift):
    return (attrgetter(shift) for (header, attrgetter) in HEADERS)


def shifts_to_csv(shifts, filename=None):
    if filename is None:
        filename = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S.csv")
    with open(filename, 'w') as f:
        writer = UnicodeWriter(f)
        writer.writerow([v[0] for v in HEADERS])
        for shift in shifts:
            writer.writerow(shift_to_row(shift))

    return filename


def all_shifts_to_csv():
    shifts_to_csv(Shift.objects.filter(owner__isnull=False).order_by(*ORDER_BY), 'all_shifts.csv')


def users_to_csv():
    for user in User.objects.filter(shifts__isnull=False):
        path = os.path.join('shifts_by_user', "{0}.csv".format(user.email))
        shifts_to_csv(user.shifts.order_by(*ORDER_BY), filename=path)


def departments_to_csv():
    for department in Department.objects.filter(shifts__isnull=False):
        path = os.path.join('shifts_by_department', "{0}-{1}.csv".format(department.pk, department.name))
        shifts_to_csv(department.shifts.filter(owner__isnull=False).order_by(*ORDER_BY), filename=path)