diff --git a/reports/graphs/sleep_pattern.py b/reports/graphs/sleep_pattern.py index 15fbf3fa..73a523ef 100644 --- a/reports/graphs/sleep_pattern.py +++ b/reports/graphs/sleep_pattern.py @@ -11,6 +11,8 @@ from core.utils import duration_string from reports import utils +from datetime import timedelta + ASLEEP_COLOR = 'rgb(35, 110, 150)' AWAKE_COLOR = 'rgba(255, 255, 255, 0)' @@ -21,9 +23,11 @@ def sleep_pattern(sleeps): :param sleeps: a QuerySet of Sleep instances. :returns: a tuple of the the graph's html and javascript. """ - days = {} last_end_time = None adjustment = None + + days = _init_days(sleeps.first().start, sleeps.last().end) + for sleep in sleeps: start_time = timezone.localtime(sleep.start) end_time = timezone.localtime(sleep.end) @@ -31,13 +35,6 @@ def sleep_pattern(sleeps): end_date = end_time.date().isoformat() duration = sleep.duration - # Ensure that lists are initialized for the start and end date (as they - # may be different dates). - if start_date not in days: - days[start_date] = [] - if end_date not in days: - days[end_date] = [] - # Check if the previous entry crossed midnight (see below). if adjustment: _add_adjustment(adjustment, days) @@ -171,6 +168,12 @@ def sleep_pattern(sleeps): return utils.split_graph_output(output) +def _init_days(first_day, last_day): + period = (last_day.date() - first_day.date()).days + 1 + def new_day(d): return (first_day + timedelta(days=d)).date().isoformat() + return {new_day(day): [] for day in range(period)} + + def _add_adjustment(adjustment, days): """ Adds "adjustment" data for entries that cross midnight.