From 6ba6102e3c73f6c19198bcaf09348b31806d0d7f Mon Sep 17 00:00:00 2001
From: Caspar Wylie <casparwylie@Caspars-MacBook-Pro.local>
Date: Fri, 3 Feb 2023 14:50:46 +0000
Subject: [PATCH] Support partial function inspection in data migrations

---
 CHANGELOG.md                                  |  1 +
 django_migration_linter/migration_linter.py   | 19 ++++++++++----
 tests/functional/test_data_migrations.py      |  9 +++++++
 .../migrations/0004_partial_function.py       | 26 +++++++++++++++++++
 4 files changed, 50 insertions(+), 5 deletions(-)
 create mode 100644 tests/test_project/app_data_migrations/migrations/0004_partial_function.py

diff --git a/CHANGELOG.md b/CHANGELOG.md
index beafd651..bcf7b2fb 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -9,6 +9,7 @@ Features:
 
 - Fixed `RunPython` model import check when using a `through` object like `MyModel.many_to_many.through.objects.filter(...)` (issue #218)
 - Mark the `IgnoreMigration` operation as `elidable=True`
+- Handle `functools.partial` functions in RunPython data migrations
 
 Bug:
 
diff --git a/django_migration_linter/migration_linter.py b/django_migration_linter/migration_linter.py
index f73636ad..5ba96053 100644
--- a/django_migration_linter/migration_linter.py
+++ b/django_migration_linter/migration_linter.py
@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import functools
 import hashlib
 import inspect
 import logging
@@ -473,10 +474,16 @@ def analyse_data_migration(
 
         return errors, ignored, warnings
 
+    @staticmethod
+    def discover_function(function):
+        if isinstance(function, functools.partial):
+            return function.func
+        return function
+
     def lint_runpython(
         self, runpython: RunPython
     ) -> tuple[list[Issue], list[Issue], list[Issue]]:
-        function_name = runpython.code.__name__
+        function_name = self.discover_function(runpython.code).__name__
         error = []
         ignored = []
         warning = []
@@ -551,8 +558,9 @@ def lint_runpython(
     def get_runpython_model_import_issues(code: Callable) -> list[Issue]:
         model_object_regex = re.compile(r"[^a-zA-Z0-9._]?([a-zA-Z0-9._]+?)\.objects")
 
-        function_name = code.__name__
-        source_code = inspect.getsource(code)
+        function = MigrationLinter.discover_function(code)
+        function_name = function.__name__
+        source_code = inspect.getsource(function)
 
         called_models = model_object_regex.findall(source_code)
         issues = []
@@ -582,8 +590,9 @@ def get_runpython_model_import_issues(code: Callable) -> list[Issue]:
     def get_runpython_model_variable_naming_issues(code: Callable) -> list[Issue]:
         model_object_regex = re.compile(r"[^a-zA-Z]?([a-zA-Z0-9]+?)\.objects")
 
-        function_name = code.__name__
-        source_code = inspect.getsource(code)
+        function = MigrationLinter.discover_function(code)
+        function_name = function.__name__
+        source_code = inspect.getsource(function)
 
         called_models = model_object_regex.findall(source_code)
         issues = []
diff --git a/tests/functional/test_data_migrations.py b/tests/functional/test_data_migrations.py
index 98ef8da7..e3c9184b 100644
--- a/tests/functional/test_data_migrations.py
+++ b/tests/functional/test_data_migrations.py
@@ -104,6 +104,15 @@ def test_warnings_as_errors_tests_no_match(self):
         self.assertEqual(0, self.linter.nb_erroneous)
         self.assertFalse(self.linter.has_errors)
 
+    def test_partial_function(self):
+        reverse_migration = self.linter.migration_loader.disk_migrations[
+            ("app_data_migrations", "0004_partial_function")
+        ]
+        self.linter.lint_migration(reverse_migration)
+
+        self.assertEqual(1, self.linter.nb_warnings)
+        self.assertFalse(self.linter.has_errors)
+
 
 class DataMigrationModelImportTestCase(unittest.TestCase):
     def test_missing_get_model_import(self):
diff --git a/tests/test_project/app_data_migrations/migrations/0004_partial_function.py b/tests/test_project/app_data_migrations/migrations/0004_partial_function.py
new file mode 100644
index 00000000..d32bdeca
--- /dev/null
+++ b/tests/test_project/app_data_migrations/migrations/0004_partial_function.py
@@ -0,0 +1,26 @@
+# Generated by Django 2.2 on 2020-01-27 17:47
+
+from __future__ import annotations
+
+import functools
+
+from django.db import migrations
+
+
+def backward_op(c, david, test):
+    pass
+
+
+def forward_op(const, arg1, arg2):
+    pass
+
+
+class Migration(migrations.Migration):
+    dependencies = [("app_data_migrations", "0003_incorrect_arguments")]
+
+    operations = [
+        migrations.RunPython(
+            functools.partial(forward_op, "constant"),
+            functools.partial(backward_op, "constant"),
+        )
+    ]