Added the templatetag.

4 files changed, 109 insertions, 1 deletions

diff --git a/src/shorturls/templatetags/__init__.py b/src/shorturls/templatetags/__init__.py
new file mode 100644
index 0000000..e69de29
--- /dev/null
+++ b/src/shorturls/templatetags/__init__.py
diff --git a/src/shorturls/templatetags/shorturl.py b/src/shorturls/templatetags/shorturl.py
new file mode 100644
index 0000000..5d4ec5e
--- /dev/null
+++ b/src/shorturls/templatetags/shorturl.py
@@ -0,0 +1,59 @@
+import urlparse
+from django import template
+from django.conf import settings
+from django.core import urlresolvers
+from django.utils.safestring import mark_safe
+from shorturls.baseconv import base62
+
+class ShortURL(template.Node):
+    @classmethod
+    def parse(cls, parser, token):
+        parts = token.split_contents()
+        if len(parts) != 2:
+            raise template.TemplateSyntaxError("%s takes exactly one argument" % parts[0])
+        return cls(template.Variable(parts[1]))
+        
+    def __init__(self, obj):
+        self.obj = obj
+        
+    def render(self, context):
+        try:
+            obj = self.obj.resolve(context)
+        except template.VariableDoesNotExist:
+            return ''
+            
+        try:
+            prefix = self.get_prefix(obj)
+        except (AttributeError, KeyError):
+            return ''
+    
+        try:
+            url = urlresolvers.reverse('shorturls.views.redirect', kwargs = {
+                'prefix': prefix,
+                'tiny': base62.from_decimal(obj.pk)
+            })
+        except urlresolvers.NoReverseMatch:
+            return ''
+        
+        if hasattr(settings, 'SHORT_BASE_URL') and settings.SHORT_BASE_URL:
+            return urlparse.urljoin(settings.SHORT_BASE_URL, url)
+        else:
+            return url
+            
+    def get_prefix(self, model):
+        if not hasattr(self.__class__, '_prefixmap'):
+            self.__class__._prefixmap = dict((m,p) for p,m in settings.SHORTEN_MODELS.items())
+        key = '%s.%s' % (model._meta.app_label, model.__class__.__name__.lower())
+        return self.__class__._prefixmap[key]
+        
+class RevCanonical(ShortURL):
+    def render(self, context):
+        url = super(RevCanonical, self).render(context)
+        if url:
+            return mark_safe('<link rev="canonical" href="%s">' % url)
+        else:
+            return ''
+
+register = template.Library()
+register.tag('shorturl', ShortURL.parse)
+register.tag('revcanonical', RevCanonical.parse)
\ No newline at end of file
diff --git a/src/shorturls/tests/__init__.py b/src/shorturls/tests/__init__.py
index 6b7a7a3..25d0b61 100644
--- a/src/shorturls/tests/__init__.py
+++ b/src/shorturls/tests/__init__.py
@@ -1,2 +1,3 @@
 from shorturls.tests.models import *
-from shorturls.tests.test_views import *
\ No newline at end of file
+from shorturls.tests.test_views import *
+from shorturls.tests.test_templatetag import *
\ No newline at end of file
diff --git a/src/shorturls/tests/test_templatetag.py b/src/shorturls/tests/test_templatetag.py
new file mode 100644
index 0000000..2feef91
--- /dev/null
+++ b/src/shorturls/tests/test_templatetag.py
@@ -0,0 +1,48 @@
+from django import template
+from django.conf import settings
+from django.test import TestCase
+from shorturls.tests.models import Animal, Vegetable, Mineral
+
+class RedirectViewTestCase(TestCase):
+    urls = 'shorturls.urls'
+    fixtures = ['shorturls-test-data.json']
+
+    def setUp(self):
+        self.old_shorten = getattr(settings, 'SHORTEN_MODELS', None)
+        self.old_base = getattr(settings, 'SHORT_BASE_URL', None)
+        settings.SHORT_BASE_URL = None
+        settings.SHORTEN_MODELS = {
+            'A': 'shorturls.animal',
+            'V': 'shorturls.vegetable',
+        }
+        
+    def tearDown(self):
+        if self.old_shorten is not None:
+            settings.SHORTEN_MODELS = self.old_shorten
+        if self.old_base is not None:
+            settings.SHORT_BASE_URL = self.old_base
+
+    def render(self, t, **c):
+        return template.Template('{% load shorturl %}'+t).render(c)
+        
+    def test_shorturl(self):
+        r = self.render('{% shorturl a %}', a=Animal.objects.get(id=12345))
+        self.assertEqual(r, '/ADNH')
+        
+    def test_bad_context(self):
+        r = self.render('{% shorturl a %}')
+        self.assertEqual(r, '')
+
+    def test_no_prefix(self):
+        r = self.render('{% shorturl m %}', m=Mineral.objects.all()[0])
+        self.assertEqual(r, '')
+        
+    def test_short_base_url(self):
+        settings.SHORT_BASE_URL = 'http://example.com/'
+        r = self.render('{% shorturl a %}', a=Animal.objects.get(id=12345))
+        self.assertEqual(r, 'http://example.com/ADNH')
+        
+    def test_revcanonical(self):
+        r = self.render('{% revcanonical a %}', a=Animal.objects.get(id=12345))
+        self.assertEqual(r, '<link rev="canonical" href="/ADNH">')
+        
\ No newline at end of file