Epidemic Week Calculator

epiweek.py

from datetime import date
from datetime import timedelta
import copy

# ported from npmjs epi-week package
# https://github.com/wombleton/epi-week

#
#getFirstWeek = (year) ->
#  end = new Date(year, 0, 1)
#  end.setHours(0, 0, 0, 0)
#  days = 0
#  while ++days < 4 or end.getDay() isnt 6
#    end.setDate(end.getDate() + 1)
#  start = new Date(end.getTime())
#  start.setDate(start.getDate() - 6)
#  start

def getFirstWeek(year):
    """
        >>> getFirstWeek(2012).strftime('%A %Y-%m-%d')
        'Sunday 2012-01-01'
        >>> getFirstWeek(2011).strftime('%A %Y-%m-%d')
        'Sunday 2011-01-02'
        >>> getFirstWeek(2013).strftime('%A %Y-%m-%d')
        'Sunday 2012-12-30'
        >>> getFirstWeek(2007).strftime('%A %Y-%m-%d')
        'Sunday 2006-12-31'

    """
    end = date(year, 1, 1)
    days = 0
    while days < 3 or end.weekday() != 5:
        end += timedelta(1)
        days += 1
    start = end - timedelta(6)
    return start
        
    
#
#calculate = (date = new Date(), marker) ->
#  date.setHours(0, 0, 0, 0)
#
#  firstWeekThisYear = getFirstWeek(date.getFullYear())
#  firstWeekNextYear = getFirstWeek(date.getFullYear() + 1)
#
#  marker ?= new Date(date.getTime())
#  marker.setMonth(0, 1)
#  week = 0
#  while marker < date
#    day = marker.getDay()
#    week++ if day is 6
#    marker.setDate(marker.getDate() + 1)
#  if week is 0
#    if date >= firstWeekThisYear
#      week: 1
#      year: date.getFullYear()
#    else
#      marker.setFullYear(date.getFullYear() - 1, 0, 1)
#      calculate(date, marker)
#  else
#    if week >= 51 and date >= firstWeekNextYear
#      week: 1
#      year: date.getFullYear() + 1
#    else
#      week: week + 1
#      year: date.getFullYear() - (if date < firstWeekThisYear then 1 else 0)

def calculate(d, marker=None):
    """
        >>> calculate(date(2013,4,24))
        (17, 2013)
        >>> calculate(date(2013,4,28))
        (18, 2013)
        >>> calculate(date(2013,4,27))
        (17, 2013)
        >>> calculate(date(2013,4,21))
        (17, 2013)
        >>> calculate(date(2013,4,20))
        (16, 2013)
        >>> calculate(date(2013,12,31))
        (1, 2014)
        >>> calculate(date(2007,1,1))
        (1, 2007)
        >>> calculate(date(2006,12,30))
        (52, 2006)
    """
    firstWeekThisYear = getFirstWeek(d.year)
    firstWeekNextYear = getFirstWeek(d.year + 1)
    if marker is None:
        marker = date(d.year, 1, 1)
    else:
        marker = date(marker.year, 1, 1)
    week = 0
    while marker < d:
        day = marker.weekday()
        if day == 5:
            week += 1
        marker += timedelta(1)
    if week == 0:
        if d >= firstWeekThisYear:
            return (1, d.year)
        else:
            marker = date(marker.year -1, 1, 1)
            return calculate(d, marker)
    else:
        if week >= 51 and d >= firstWeekNextYear:
            return (1, d.year + 1)
        else:
            return (week + 1, d.year - (1 if d < firstWeekThisYear else 0))

def first_day(week, year):
    dt = getFirstWeek(year)
    for i in range(week - 1):
        dt += timedelta(7)
    return dt

def test_calculate():
    for d, o in [
        ((2014,1,1), (1,2014)),
        ((2013,12,31), (1,2014)),
        ((2007,1,1), (1,2007)),
        ((2006,12,31), (1,2007)),
        ((2006,12,30), (52,2006))
    ]:
        assert epiweek.calculate(date(*d)) == o