From 5d4d5ca5c13cd38a1c00c7adf471f57bdc12d6d5 Mon Sep 17 00:00:00 2001
From: Adam Johnson <me@adamj.eu>
Date: Sun, 24 Sep 2023 21:01:08 +0100
Subject: [PATCH] Fix admin.register custom site ordering bug (#338)

---
 CHANGELOG.rst                               |  4 ++
 src/django_upgrade/fixers/admin_register.py | 53 +++++++++++++++++++--
 tests/fixers/test_admin_register.py         | 51 ++++++++++++++++++++
 3 files changed, 105 insertions(+), 3 deletions(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index d4cc39f7..5641e532 100644
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -12,6 +12,10 @@ Changelog
 
   Thanks to Thibaut Decombe in `PR #368 <https://github.com/adamchainz/django-upgrade/pull/368>`__.
 
+* Fix issue with ``@admin.register()`` checker
+
+  Thanks to Jan Pieter Waagmeester for the report in `Issue #337 <https://github.com/adamchainz/django-upgrade/issues/337>`__, and to Thibaut Decombe for the review in `PR #338 <https://github.com/adamchainz/django-upgrade/pull/338>`__.
+
 1.14.1 (2023-08-16)
 -------------------
 
diff --git a/src/django_upgrade/fixers/admin_register.py b/src/django_upgrade/fixers/admin_register.py
index 7093fb93..eb0ab053 100644
--- a/src/django_upgrade/fixers/admin_register.py
+++ b/src/django_upgrade/fixers/admin_register.py
@@ -36,10 +36,11 @@
 
 
 class AdminDetails:
-    __slots__ = ("parent", "model_names_per_site")
+    __slots__ = ("parent", "lineno", "model_names_per_site")
 
-    def __init__(self, parent: ast.AST) -> None:
+    def __init__(self, parent: ast.AST, lineno: int) -> None:
         self.parent = parent
+        self.lineno = lineno
         self.model_names_per_site: dict[str, set[str]] = {}
 
 
@@ -65,7 +66,9 @@ def visit_ClassDef(
     parents: list[ast.AST],
 ) -> Iterable[tuple[Offset, TokenFunc]]:
     if _is_django_admin_imported(state) and not uses_full_super_in_init_or_new(node):
-        decorable_admins.setdefault(state, {})[node.name] = AdminDetails(parents[-1])
+        decorable_admins.setdefault(state, {})[node.name] = AdminDetails(
+            parents[-1], node.lineno
+        )
         if not node.decorator_list:
             offset = ast_start_offset(node)
             decorated = False
@@ -201,6 +204,18 @@ def visit_Call(
                 and admin_details is not None
                 and admin_details.parent == parents[-2]
                 and not (site_name and not admin_name.endswith("Admin"))
+                and (
+                    site_name == ""
+                    or (
+                        (
+                            site_defined_line := get_site_defined_line(
+                                parents[0], site_name
+                            )
+                        )
+                        is not None
+                        and site_defined_line < admin_details.lineno
+                    )
+                )
             ):
                 admin_details.model_names_per_site.setdefault(site_name, set()).update(
                     model_names
@@ -248,3 +263,35 @@ def visit_Call(
                     state_details[site_name] = unregistered_names
                 elif existing_names is not True:
                     existing_names.update(unregistered_names)
+
+
+site_definitions: MutableMapping[
+    ast.Module, dict[str, int | None]
+] = WeakKeyDictionary()
+
+
+def get_site_defined_line(module: ast.AST, site_name: str) -> int | None:
+    assert isinstance(module, ast.Module)
+    lines = site_definitions.get(module, None)
+    if lines is None:
+        lines = {}
+        for node in module.body:
+            if isinstance(node, ast.ImportFrom):
+                for alias in node.names:
+                    if alias.asname is not None:
+                        name = alias.asname
+                    else:
+                        name = alias.name
+
+                    if name.endswith("site") and name not in lines:
+                        lines[name] = node.lineno
+            elif (
+                isinstance(node, ast.Assign)
+                and len(node.targets) == 1
+                and isinstance(node.targets[0], ast.Name)
+                and (name := node.targets[0].id) not in lines
+            ):
+                lines[name] = node.lineno
+
+        site_definitions[module] = lines
+    return lines.get(site_name, None)
diff --git a/tests/fixers/test_admin_register.py b/tests/fixers/test_admin_register.py
index a68ed1e0..39234289 100644
--- a/tests/fixers/test_admin_register.py
+++ b/tests/fixers/test_admin_register.py
@@ -782,6 +782,57 @@ class MyModelAdmin(CustomModelAdmin):
     )
 
 
+def test_custom_admin_site_defined_after_admin():
+    check_noop(
+        """\
+        from django.contrib import admin
+
+        class MyModelAdmin(admin.ModelAdmin):
+            pass
+
+        custom_site = admin.AdminSite(...)
+
+        custom_site.register(MyModel, MyModelAdmin)
+        """,
+        settings=settings,
+        filename="admin.py",
+    )
+
+
+def test_custom_admin_site_defined_after_admin_import():
+    check_noop(
+        """\
+        from django.contrib import admin
+
+        class MyModelAdmin(admin.ModelAdmin):
+            pass
+
+        from myapp.admin import custom_site
+
+        custom_site.register(MyModel, MyModelAdmin)
+        """,
+        settings=settings,
+        filename="admin.py",
+    )
+
+
+def test_custom_admin_site_defined_after_admin_import_as():
+    check_noop(
+        """\
+        from django.contrib import admin
+
+        class MyModelAdmin(admin.ModelAdmin):
+            pass
+
+        from myapp.admin import site as custom_site
+
+        custom_site.register(MyModel, MyModelAdmin)
+        """,
+        settings=settings,
+        filename="admin.py",
+    )
+
+
 def test_custom_admin_site():
     check_transformed(
         """\