Ticket #10821: if_in_list_tag.py

# This tag library adds a useful tag for checking whether an element is in a list of not
# This can be useful if iterating over a list of all values, and testing to see if it is
# in our current collection. It behaves like a usual 'if' and can be accompanied by an 
# 'else' block as well.
#
# EG:
#   ctxt = { 'my_colours': ['orange', 'yellow'], 
#            'all_colours': ['red', 'orange', 'yellow', 'green', blue', 'indigo', 'violet'] }
# <tr>{% for clr in all_colours %}<th>{{ clr }}</th>{% endfor %}</tr>
# <tr>
#   {% for clr in all_colours %}
#   <td>{% IfInList clr in my_colours %}I like it{% else %}I hate it{% endif %}</td>
#   {% endfor %}
# </tr>
#
# You can also negate the test:
#
#  <td>{% IfInList clr not in mine %}</td>
#
from django.template import Node, Variable, TemplateSyntaxError, Library, NodeList, VariableDoesNotExist

register = Library()

class IfInListNode(Node):
  def __init__(self, obj, list, nodelist_true, nodelist_false, negate):
    self.obj, self.list = Variable(obj), list
    self.nodelist_true, self.nodelist_false = nodelist_true, nodelist_false
    self.negate = negate

  def __repr__(self):
    return '<IfInList node>'

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

  def get_nodes_by_type(self, nodetype):
    nodes = []
    if isinstance(self, nodetype):
      nodes.append(self)
    nodes.extend(self.nodelist_true.get_nodes_by_type(nodetype))
    nodes.extend(self.nodelist_false.get_nodes_by_type(nodetype))
    return nodes

  def render(self, context):
    try:
      val = self.obj.resolve(context)
    except VariableDoesNotExist:
      val = None
    try:
      list = self.list.resolve(context, True)
    except VariableDoesNotExist:
      list = []
    in_list = val in list
    if in_list and not self.negate or self.negate and not in_list:
      return self.nodelist_true.render(context)
    else:
      return self.nodelist_false.render(context)


@register.tag(name='IfInList')
def do_if_in_list(parser, token):
  bits = token.contents.split()
  tag_name, var_not, var_in = None, None, None
  if len(bits) == 4:
    tag_name, obj, var_in, list = token.split_contents()
  elif len(bits) == 5:
    tag_name, obj, var_not, var_in, list = token.split_contents()
  if not tag_name  \
      or (var_not is not None and var_not != 'not') \
      or (var_in is not None and var_in != 'in'):
        raise TemplateSyntaxError("Invalid format! 'IfInList obj [not] in list'. You had: %s" % token.contents)
  negate = var_not is not None
  list = parser.compile_filter(list)
  nodelist_true = parser.parse(('else', 'endif'))
  tok = parser.next_token()
  if tok.contents == 'else':
    nodelist_false = parser.parse(('endif',))
    parser.delete_first_token()
  else:
    nodelist_false = NodeList()
  return IfInListNode(obj, list, nodelist_true, nodelist_false, negate)