Fixed #24629 -- Unified Transform and Expression APIs

2 files changed, 157 insertions, 8 deletions

diff --git a/tests/custom_lookups/tests.py b/tests/custom_lookups/tests.py
index d8852e1e3c..c538d23b76 100644
--- a/tests/custom_lookups/tests.py
+++ b/tests/custom_lookups/tests.py
@@ -126,11 +126,17 @@ class YearLte(models.lookups.LessThanOrEqual):
         return "%s <= (%s || '-12-31')::date" % (lhs_sql, rhs_sql), params
 
 
-class SQLFunc(models.Lookup):
-    def __init__(self, name, *args, **kwargs):
-        super(SQLFunc, self).__init__(*args, **kwargs)
-        self.name = name
+class Exactly(models.lookups.Exact):
+    """
+    This lookup is used to test lookup registration.
+    """
+    lookup_name = 'exactly'
 
+    def get_rhs_op(self, connection, rhs):
+        return connection.operators['exact'] % rhs
+
+
+class SQLFuncMixin(object):
     def as_sql(self, compiler, connection):
         return '%s()', [self.name]
 
@@ -139,13 +145,28 @@ class SQLFunc(models.Lookup):
         return CustomField()
 
 
+class SQLFuncLookup(SQLFuncMixin, models.Lookup):
+    def __init__(self, name, *args, **kwargs):
+        super(SQLFuncLookup, self).__init__(*args, **kwargs)
+        self.name = name
+
+
+class SQLFuncTransform(SQLFuncMixin, models.Transform):
+    def __init__(self, name, *args, **kwargs):
+        super(SQLFuncTransform, self).__init__(*args, **kwargs)
+        self.name = name
+
+
 class SQLFuncFactory(object):
 
-    def __init__(self, name):
+    def __init__(self, key, name):
+        self.key = key
         self.name = name
 
     def __call__(self, *args, **kwargs):
-        return SQLFunc(self.name, *args, **kwargs)
+        if self.key == 'lookupfunc':
+            return SQLFuncLookup(self.name, *args, **kwargs)
+        return SQLFuncTransform(self.name, *args, **kwargs)
 
 
 class CustomField(models.TextField):
@@ -153,13 +174,13 @@ class CustomField(models.TextField):
     def get_lookup(self, lookup_name):
         if lookup_name.startswith('lookupfunc_'):
             key, name = lookup_name.split('_', 1)
-            return SQLFuncFactory(name)
+            return SQLFuncFactory(key, name)
         return super(CustomField, self).get_lookup(lookup_name)
 
     def get_transform(self, lookup_name):
         if lookup_name.startswith('transformfunc_'):
             key, name = lookup_name.split('_', 1)
-            return SQLFuncFactory(name)
+            return SQLFuncFactory(key, name)
         return super(CustomField, self).get_transform(lookup_name)
 
 
@@ -200,6 +221,27 @@ class DateTimeTransform(models.Transform):
 
 
 class LookupTests(TestCase):
+
+    def test_custom_name_lookup(self):
+        a1 = Author.objects.create(name='a1', birthdate=date(1981, 2, 16))
+        Author.objects.create(name='a2', birthdate=date(2012, 2, 29))
+        custom_lookup_name = 'isactually'
+        custom_transform_name = 'justtheyear'
+        try:
+            models.DateField.register_lookup(YearTransform)
+            models.DateField.register_lookup(YearTransform, custom_transform_name)
+            YearTransform.register_lookup(Exactly)
+            YearTransform.register_lookup(Exactly, custom_lookup_name)
+            qs1 = Author.objects.filter(birthdate__testyear__exactly=1981)
+            qs2 = Author.objects.filter(birthdate__justtheyear__isactually=1981)
+            self.assertQuerysetEqual(qs1, [a1], lambda x: x)
+            self.assertQuerysetEqual(qs2, [a1], lambda x: x)
+        finally:
+            YearTransform._unregister_lookup(Exactly)
+            YearTransform._unregister_lookup(Exactly, custom_lookup_name)
+            models.DateField._unregister_lookup(YearTransform)
+            models.DateField._unregister_lookup(YearTransform, custom_transform_name)
+
     def test_basic_lookup(self):
         a1 = Author.objects.create(name='a1', age=1)
         a2 = Author.objects.create(name='a2', age=2)
@@ -299,6 +341,19 @@ class BilateralTransformTests(TestCase):
             with self.assertRaises(NotImplementedError):
                 Author.objects.filter(name__upper__in=Author.objects.values_list('name'))
 
+    def test_bilateral_multi_value(self):
+        with register_lookup(models.CharField, UpperBilateralTransform):
+            Author.objects.bulk_create([
+                Author(name='Foo'),
+                Author(name='Bar'),
+                Author(name='Ray'),
+            ])
+            self.assertQuerysetEqual(
+                Author.objects.filter(name__upper__in=['foo', 'bar', 'doe']).order_by('name'),
+                ['Bar', 'Foo'],
+                lambda a: a.name
+            )
+
     def test_div3_bilateral_extract(self):
         with register_lookup(models.IntegerField, Div3BilateralTransform):
             a1 = Author.objects.create(name='a1', age=1)
diff --git a/tests/db_functions/tests.py b/tests/db_functions/tests.py
index a401c550fb..00aac82a7b 100644
--- a/tests/db_functions/tests.py
+++ b/tests/db_functions/tests.py
@@ -547,3 +547,97 @@ class FunctionTests(TestCase):
             ['How to Time Travel'],
             lambda a: a.title
         )
+
+    def test_length_transform(self):
+        try:
+            CharField.register_lookup(Length, 'length')
+            Author.objects.create(name='John Smith', alias='smithj')
+            Author.objects.create(name='Rhonda')
+            authors = Author.objects.filter(name__length__gt=7)
+            self.assertQuerysetEqual(
+                authors.order_by('name'), [
+                    'John Smith',
+                ],
+                lambda a: a.name
+            )
+        finally:
+            CharField._unregister_lookup(Length, 'length')
+
+    def test_lower_transform(self):
+        try:
+            CharField.register_lookup(Lower, 'lower')
+            Author.objects.create(name='John Smith', alias='smithj')
+            Author.objects.create(name='Rhonda')
+            authors = Author.objects.filter(name__lower__exact='john smith')
+            self.assertQuerysetEqual(
+                authors.order_by('name'), [
+                    'John Smith',
+                ],
+                lambda a: a.name
+            )
+        finally:
+            CharField._unregister_lookup(Lower, 'lower')
+
+    def test_upper_transform(self):
+        try:
+            CharField.register_lookup(Upper, 'upper')
+            Author.objects.create(name='John Smith', alias='smithj')
+            Author.objects.create(name='Rhonda')
+            authors = Author.objects.filter(name__upper__exact='JOHN SMITH')
+            self.assertQuerysetEqual(
+                authors.order_by('name'), [
+                    'John Smith',
+                ],
+                lambda a: a.name
+            )
+        finally:
+            CharField._unregister_lookup(Upper, 'upper')
+
+    def test_func_transform_bilateral(self):
+        class UpperBilateral(Upper):
+            bilateral = True
+
+        try:
+            CharField.register_lookup(UpperBilateral, 'upper')
+            Author.objects.create(name='John Smith', alias='smithj')
+            Author.objects.create(name='Rhonda')
+            authors = Author.objects.filter(name__upper__exact='john smith')
+            self.assertQuerysetEqual(
+                authors.order_by('name'), [
+                    'John Smith',
+                ],
+                lambda a: a.name
+            )
+        finally:
+            CharField._unregister_lookup(UpperBilateral, 'upper')
+
+    def test_func_transform_bilateral_multivalue(self):
+        class UpperBilateral(Upper):
+            bilateral = True
+
+        try:
+            CharField.register_lookup(UpperBilateral, 'upper')
+            Author.objects.create(name='John Smith', alias='smithj')
+            Author.objects.create(name='Rhonda')
+            authors = Author.objects.filter(name__upper__in=['john smith', 'rhonda'])
+            self.assertQuerysetEqual(
+                authors.order_by('name'), [
+                    'John Smith',
+                    'Rhonda',
+                ],
+                lambda a: a.name
+            )
+        finally:
+            CharField._unregister_lookup(UpperBilateral, 'upper')
+
+    def test_function_as_filter(self):
+        Author.objects.create(name='John Smith', alias='SMITHJ')
+        Author.objects.create(name='Rhonda')
+        self.assertQuerysetEqual(
+            Author.objects.filter(alias=Upper(V('smithj'))),
+            ['John Smith'], lambda x: x.name
+        )
+        self.assertQuerysetEqual(
+            Author.objects.exclude(alias=Upper(V('smithj'))),
+            ['Rhonda'], lambda x: x.name
+        )