Replaced try..except blocks by context manager in custom lookups tests

tests/custom_lookups/tests.py

 from __future__ import unicode_literals
 
+import contextlib
 import time
 import unittest
 from datetime import date, datetime
@@ -12,6 +13,17 @@ from django.utils import timezone
 from .models import Author, MySQLUnixTimestamp
 
 
+@contextlib.contextmanager
+def register_lookup(field, *lookups):
+    try:
+        for lookup in lookups:
+            field.register_lookup(lookup)
+        yield
+    finally:
+        for lookup in lookups:
+            field._unregister_lookup(lookup)
+
+
 class Div3Lookup(models.Lookup):
     lookup_name = 'div3'
 
@@ -193,8 +205,7 @@ class LookupTests(TestCase):
         a2 = Author.objects.create(name='a2', age=2)
         a3 = Author.objects.create(name='a3', age=3)
         a4 = Author.objects.create(name='a4', age=4)
-        models.IntegerField.register_lookup(Div3Lookup)
-        try:
+        with register_lookup(models.IntegerField, Div3Lookup):
             self.assertQuerysetEqual(
                 Author.objects.filter(age__div3=0),
                 [a3], lambda x: x
@@ -211,8 +222,6 @@ class LookupTests(TestCase):
                 Author.objects.filter(age__div3=3),
                 [], lambda x: x
             )
-        finally:
-            models.IntegerField._unregister_lookup(Div3Lookup)
 
     @unittest.skipUnless(connection.vendor == 'postgresql', "PostgreSQL specific SQL used")
     def test_birthdate_month(self):
@@ -220,8 +229,7 @@ class LookupTests(TestCase):
         a2 = Author.objects.create(name='a2', birthdate=date(2012, 2, 29))
         a3 = Author.objects.create(name='a3', birthdate=date(2012, 1, 31))
         a4 = Author.objects.create(name='a4', birthdate=date(2012, 3, 1))
-        models.DateField.register_lookup(InMonth)
-        try:
+        with register_lookup(models.DateField, InMonth):
             self.assertQuerysetEqual(
                 Author.objects.filter(birthdate__inmonth=date(2012, 1, 15)),
                 [a3], lambda x: x
@@ -242,12 +250,9 @@ class LookupTests(TestCase):
                 Author.objects.filter(birthdate__inmonth=date(2012, 4, 1)),
                 [], lambda x: x
             )
-        finally:
-            models.DateField._unregister_lookup(InMonth)
 
     def test_div3_extract(self):
-        models.IntegerField.register_lookup(Div3Transform)
-        try:
+        with register_lookup(models.IntegerField, Div3Transform):
             a1 = Author.objects.create(name='a1', age=1)
             a2 = Author.objects.create(name='a2', age=2)
             a3 = Author.objects.create(name='a3', age=3)
@@ -271,15 +276,12 @@ class LookupTests(TestCase):
             self.assertQuerysetEqual(
                 baseqs.filter(age__div3__range=(1, 2)),
                 [a1, a2, a4], lambda x: x)
-        finally:
-            models.IntegerField._unregister_lookup(Div3Transform)
 
 
 class BilateralTransformTests(TestCase):
 
     def test_bilateral_upper(self):
-        models.CharField.register_lookup(UpperBilateralTransform)
-        try:
+        with register_lookup(models.CharField, UpperBilateralTransform):
             Author.objects.bulk_create([
                 Author(name='Doe'),
                 Author(name='doe'),
@@ -291,20 +293,14 @@ class BilateralTransformTests(TestCase):
             self.assertQuerysetEqual(
                 Author.objects.filter(name__upper__contains='f'),
                 ["<Author: Foo>"], ordered=False)
-        finally:
-            models.CharField._unregister_lookup(UpperBilateralTransform)
 
     def test_bilateral_inner_qs(self):
-        models.CharField.register_lookup(UpperBilateralTransform)
-        try:
+        with register_lookup(models.CharField, UpperBilateralTransform):
             with self.assertRaises(NotImplementedError):
                 Author.objects.filter(name__upper__in=Author.objects.values_list('name'))
-        finally:
-            models.CharField._unregister_lookup(UpperBilateralTransform)
 
     def test_div3_bilateral_extract(self):
-        models.IntegerField.register_lookup(Div3BilateralTransform)
-        try:
+        with register_lookup(models.IntegerField, Div3BilateralTransform):
             a1 = Author.objects.create(name='a1', age=1)
             a2 = Author.objects.create(name='a2', age=2)
             a3 = Author.objects.create(name='a3', age=3)
@@ -328,13 +324,9 @@ class BilateralTransformTests(TestCase):
             self.assertQuerysetEqual(
                 baseqs.filter(age__div3__range=(1, 2)),
                 [a1, a2, a4], lambda x: x)
-        finally:
-            models.IntegerField._unregister_lookup(Div3BilateralTransform)
 
     def test_bilateral_order(self):
-        models.IntegerField.register_lookup(Mult3BilateralTransform)
-        models.IntegerField.register_lookup(Div3BilateralTransform)
-        try:
+        with register_lookup(models.IntegerField, Mult3BilateralTransform, Div3BilateralTransform):
             a1 = Author.objects.create(name='a1', age=1)
             a2 = Author.objects.create(name='a2', age=2)
             a3 = Author.objects.create(name='a3', age=3)
@@ -348,13 +340,9 @@ class BilateralTransformTests(TestCase):
             self.assertQuerysetEqual(
                 baseqs.filter(age__div3__mult3=42),
                 [a3], lambda x: x)
-        finally:
-            models.IntegerField._unregister_lookup(Mult3BilateralTransform)
-            models.IntegerField._unregister_lookup(Div3BilateralTransform)
 
     def test_bilateral_fexpr(self):
-        models.IntegerField.register_lookup(Mult3BilateralTransform)
-        try:
+        with register_lookup(models.IntegerField, Mult3BilateralTransform):
             a1 = Author.objects.create(name='a1', age=1, average_rating=3.2)
             a2 = Author.objects.create(name='a2', age=2, average_rating=0.5)
             a3 = Author.objects.create(name='a3', age=3, average_rating=1.5)
@@ -367,23 +355,18 @@ class BilateralTransformTests(TestCase):
                 # Same as age >= average_rating
                 baseqs.filter(age__mult3__gte=models.F('average_rating')),
                 [a2, a3], lambda x: x)
-        finally:
-            models.IntegerField._unregister_lookup(Mult3BilateralTransform)
 
 
 @override_settings(USE_TZ=True)
 class DateTimeLookupTests(TestCase):
     @unittest.skipUnless(connection.vendor == 'mysql', "MySQL specific SQL used")
     def test_datetime_output_field(self):
-        models.PositiveIntegerField.register_lookup(DateTimeTransform)
-        try:
+        with register_lookup(models.PositiveIntegerField, DateTimeTransform):
             ut = MySQLUnixTimestamp.objects.create(timestamp=time.time())
             y2k = timezone.make_aware(datetime(2000, 1, 1))
             self.assertQuerysetEqual(
                 MySQLUnixTimestamp.objects.filter(timestamp__as_datetime__gt=y2k),
                 [ut], lambda x: x)
-        finally:
-            models.PositiveIntegerField._unregister_lookup(DateTimeTransform)
 
 
 class YearLteTests(TestCase):
@@ -514,8 +497,7 @@ class TrackCallsYearTransform(YearTransform):
 
 class LookupTransformCallOrderTests(TestCase):
     def test_call_order(self):
-        models.DateField.register_lookup(TrackCallsYearTransform)
-        try:
+        with register_lookup(models.DateField, TrackCallsYearTransform):
             # junk lookup - tries lookup, then transform, then fails
             with self.assertRaises(FieldError):
                 Author.objects.filter(birthdate__testyear__junk=2012)
@@ -538,9 +520,6 @@ class LookupTransformCallOrderTests(TestCase):
             self.assertEqual(TrackCallsYearTransform.call_order,
                              ['lookup'])
 
-        finally:
-            models.DateField._unregister_lookup(TrackCallsYearTransform)
-
 
 class CustomisedMethodsTests(TestCase):
 
@@ -563,8 +542,7 @@ class CustomisedMethodsTests(TestCase):
 
 class SubqueryTransformTests(TestCase):
     def test_subquery_usage(self):
-        models.IntegerField.register_lookup(Div3Transform)
-        try:
+        with register_lookup(models.IntegerField, Div3Transform):
             Author.objects.create(name='a1', age=1)
             a2 = Author.objects.create(name='a2', age=2)
             Author.objects.create(name='a3', age=3)
@@ -572,5 +550,3 @@ class SubqueryTransformTests(TestCase):
             self.assertQuerysetEqual(
                 Author.objects.order_by('name').filter(id__in=Author.objects.filter(age__div3=2)),
                 [a2], lambda x: x)
-        finally:
-            models.IntegerField._unregister_lookup(Div3Transform)