Template Tags for Time Branching: ifbefore, ifafter, ifafterbefore

# -*- coding: utf-8 -*-
"""
Template tag to determine and display based on time
Example:
{% ifbefore '2011-12-24 20:00:00' %}
    This is displayed before 2011-12-24 20:00:00
{% else %}
    This is displayed after 2011-12-24 20:00:00
{% endifbefore %}

Arguments can be datetime or date strings in the formats
'%Y-%m-%d %H:%M:%S', '%Y-%m-%d', '%H:%M:%S'
"""

import datetime

from django import template
from django.template import Node, NodeList, resolve_variable

register = template.Library()

class IfAfterBeforeNode(Node):
    child_nodelists = ('nodelist_true', 'nodelist_false')
    datestr_format_list = ('%Y-%m-%d %H:%M:%S', '%Y-%m-%d', '%H:%M:%S')

    def __init__(self, var_after, var_before, nodelist_true, nodelist_false=None):
        self.nodelist_true, self.nodelist_false = nodelist_true, nodelist_false
        self.var_after = var_after
        self.var_before = var_before
        self.now = datetime.datetime.now()

    def strptime(self, date_string):
        """
        Parse date string into datetime.
        Returns None if parsing fails.
        """
        date_string = date_string.strip("""'" """)
        for datestr_format in self.datestr_format_list:
            try:
                return datetime.datetime.strptime(date_string, datestr_format)
            except ValueError:
                pass
        return None

    def __repr__(self):
        return "<IfBeforeAfter node>"

    def __iter__(self):
        for node in self.nodelist_true:
            yield node
        for node in self.nodelist_false:
            yield node

    def arg_to_dt_or_none(self, arg, context):
        """
        Retrieve argument as datetime or none
        Used for self.var_after, self.var_before
        """
        if not arg:
            return None
        try:
            # Argument may be datetime
            dt = resolve_variable(arg, context)
        except template.VariableDoesNotExist:
            dt  = None
        if dt is None or not isinstance(dt,  datetime.datetime):
            dt = self.strptime(arg)
        return dt

    def render(self, context):

        dt_after  = self.arg_to_dt_or_none(self.var_after,  context)
        dt_before = self.arg_to_dt_or_none(self.var_before, context)

        if dt_after and dt_before:
            if dt_after <= self.now < dt_before:
                return self.nodelist_true.render(context)
            else:
                return self.nodelist_false.render(context)
        elif dt_after:
            if dt_after <= self.now:
                return self.nodelist_true.render(context)
            else:
                return self.nodelist_false.render(context)
        elif dt_before:
            if self.now < dt_before:
                return self.nodelist_true.render(context)
            else:
                return self.nodelist_false.render(context)
        else:
            return u'Parse Error. No after, before'

@register.tag
def ifbefore(parser, token):
    """
    Render if current time is before "YYYY-MM-DD HH:MM:SS"
    """
    _tagname, var_before = token.split_contents()
    nodelist_true = parser.parse(('else', 'endifbefore'))
    token = parser.next_token()
    if token.contents == 'else':
        nodelist_false = parser.parse(('endifbefore',))
        parser.delete_first_token()
    else:
        nodelist_false = NodeList()
    return IfAfterBeforeNode(None, var_before, nodelist_true, nodelist_false)

@register.tag
def ifafter(parser, token):
    """
    Render if current time is after "YYYY-MM-DD HH:MM:SS"
    """
    _tagname, var_after = token.split_contents()
    nodelist_true = parser.parse(('else', 'endifafter'))
    token = parser.next_token()
    if token.contents == 'else':
        nodelist_false = parser.parse(('endifafter',))
        parser.delete_first_token()
    else:
        nodelist_false = NodeList()
    return IfAfterBeforeNode(var_after, None, nodelist_true, nodelist_false)

@register.tag
def ifafterbefore(parser, token):
    """
    Take two arguments and render if the current time is between them
    """
    _tagname, var_after, var_before = token.split_contents()

    nodelist_true = parser.parse(('else', 'endifafterbefore'))
    token = parser.next_token()
    if token.contents == 'else':
        nodelist_false = parser.parse(('endifafterbefore',))
        parser.delete_first_token()
    else:
        nodelist_false = NodeList()
    return IfAfterBeforeNode(var_after, var_before, nodelist_true, nodelist_false)