<s> from django import template <EOL> from userena . contrib . umessages . models import MessageRecipient <EOL> import re <EOL> register = template . Library ( ) <EOL> class MessageCount ( template . Node ) : <EOL> def __init__ ( self , um_from_user , var_name , um_to_user = None ) : <EOL> self . user = template . Variable ( um_from_user ) <EOL> self . var_name = var_name <EOL> if um_to_user : <EOL> self . um_to_user = template . Variable ( um_to_user ) <EOL> else : self . um_to_user = um_to_user <EOL> def render ( self , context ) : <EOL> try : <EOL> user = self . user . resolve ( context ) <EOL> except template . VariableDoesNotExist : <EOL> return '<STR_LIT>' <EOL> if not self . um_to_user : <EOL> message_count = MessageRecipient . objects . count_unread_messages_for ( user ) <EOL> else : <EOL> try : <EOL> um_to_user = self . um_to_user . resolve ( context ) <EOL> except template . VariableDoesNotExist : <EOL> return '<STR_LIT>' <EOL> message_count = MessageRecipient . objects . count_unread_messages_between ( user , <EOL> um_to_user ) <EOL> context [ self . var_name ] = message_count <EOL> return '<STR_LIT>' <EOL> @ register . tag <EOL> def get_unread_message_count_for ( parser , token ) : <EOL> """<STR_LIT>""" <EOL> try : <EOL> tag_name , arg = token . contents . split ( None , <NUM_LIT:1> ) <EOL> except ValueError : <EOL> raise template . TemplateSyntaxError ( "<STR_LIT>" % token . contents . split ( ) [ <NUM_LIT:0> ] ) <EOL> m = re . search ( r'<STR_LIT>' , arg ) <EOL> if not m : <EOL> raise template . TemplateSyntaxError ( "<STR_LIT>" % tag_name ) <EOL> user , var_name = m . groups ( ) <EOL> return MessageCount ( user , var_name ) <EOL> @ register . tag <EOL> def get_unread_message_count_between ( parser , token ) : <EOL> """<STR_LIT>""" <EOL> try : <EOL> tag_name , arg = token . contents . split ( None , <NUM_LIT:1> ) <EOL> except ValueError : <EOL> raise template . TemplateSyntaxError ( "<STR_LIT>" % token . contents . split ( ) [ <NUM_LIT:0> ] ) <EOL> m = re . search ( r'<STR_LIT>' , arg ) <EOL> if not m : <EOL> raise template . TemplateSyntaxError ( "<STR_LIT>" % tag_name ) <EOL> um_from_user , um_to_user , var_name = m . groups ( ) <EOL> return MessageCount ( um_from_user , var_name , um_to_user ) </s>
