From 4098c268ed830f7d1acbed9da65b322cb3adb469 Mon Sep 17 00:00:00 2001
From: Brian Helba <brian.helba@kitware.com>
Date: Wed, 21 Oct 2020 19:05:03 -0400
Subject: [PATCH] Annotate admin_order_field to also accept a BaseExpression

---
 CONTRIBUTORS.md                  |  2 +-
 django_admin_display/__init__.py |  5 +++--
 tests/test_decorator.py          |  3 +++
 tests/test_mypy.py               | 32 ++++++++++++++++++++++++++++----
 4 files changed, 35 insertions(+), 7 deletions(-)

diff --git a/CONTRIBUTORS.md b/CONTRIBUTORS.md
index bde0d83..23a4cb7 100644
--- a/CONTRIBUTORS.md
+++ b/CONTRIBUTORS.md
@@ -2,4 +2,4 @@
 
 * [@adamchainz](https://github.com/adamchainz/)
 * [@escaped](https://github.com/escaped/)
-
+* [@brianhelba](https://github.com/brianhelba)
diff --git a/django_admin_display/__init__.py b/django_admin_display/__init__.py
index a2f60ca..f2624e0 100644
--- a/django_admin_display/__init__.py
+++ b/django_admin_display/__init__.py
@@ -1,6 +1,7 @@
-from typing import Callable, Optional, TypeVar
+from typing import Callable, Optional, TypeVar, Union
 
 import django
+from django.db.models.expressions import BaseExpression
 
 ReturnType = TypeVar('ReturnType')
 FuncType = Callable[..., ReturnType]
@@ -8,7 +9,7 @@
 
 
 def admin_display(
-    admin_order_field: Optional[str] = None,
+    admin_order_field: Optional[Union[str, BaseExpression]] = None,
     allow_tags: Optional[bool] = None,  # deprecated in django >= 2.0
     boolean: Optional[bool] = None,
     empty_value_display: Optional[str] = None,
diff --git a/tests/test_decorator.py b/tests/test_decorator.py
index 8e89be2..2c45693 100644
--- a/tests/test_decorator.py
+++ b/tests/test_decorator.py
@@ -1,5 +1,7 @@
 import django
 import pytest
+from django.db.models import F
+from django.db.models.functions import Lower
 
 from django_admin_display import admin_display
 
@@ -9,6 +11,7 @@
 
 OPTIONS = [
     ('admin_order_field', 'radius'),
+    ('admin_order_field', Lower(F('person_name'))),
     ('boolean', True),
     ('empty_value_display', 'Undefined'),
     ('short_description', 'Is big?'),
diff --git a/tests/test_mypy.py b/tests/test_mypy.py
index c230039..ebe9b8e 100644
--- a/tests/test_mypy.py
+++ b/tests/test_mypy.py
@@ -1,21 +1,44 @@
 import pytest
+from django.db.models import F, Func
 from mypy import api
 
 from .test_decorator import OPTIONS
 
 
+def safe_repr(value):
+    if isinstance(value, Func):
+        # Django's repr for expressions does not quote string parameters
+        # Copy all expressions, wrapping their value in a "repr"
+        value = value.copy()
+        value.set_source_expressions(
+            [
+                expression.__class__(
+                    repr(
+                        expression.name
+                        if isinstance(expression, F)
+                        else expression.value
+                    )
+                )
+                for expression in value.source_expressions
+            ]
+        )
+
+    return repr(value)
+
+
 @pytest.mark.parametrize('attribute, value', OPTIONS)
 def test_failure(attribute, value):
-    value = f'"{value}"' if isinstance(value, str) else value
     code = f'''
 from django import admin
 from django.db import models
+from django.db.models import F
+from django.db.models.functions import Lower
 
 
 class SampleAdmin(admin.ModelAdmin):
     def foo(self, obj: models.Model) -> int:
         return 1
-    foo.{attribute} = {value}
+    foo.{attribute} = {safe_repr(value)}
     '''
 
     result = api.run(['-c', code])
@@ -25,15 +48,16 @@ def foo(self, obj: models.Model) -> int:
 
 @pytest.mark.parametrize('attribute, value', OPTIONS)
 def test_success(attribute, value):
-    value = f'"{value}"' if isinstance(value, str) else value
     code = f'''
 from django import admin
 from django.db import models
 from django_admin_display import admin_display
+from django.db.models import F
+from django.db.models.functions import Lower
 
 
 class SampleAdmin(admin.ModelAdmin):
-    @admin_display({attribute}={value})
+    @admin_display({attribute}={safe_repr(value)})
     def foo(self, obj: models.Model) -> int:
         return 1
     '''