Allow callables as the argument to RunPython

django/db/migrations/operations/special.py

@@ -1,6 +1,7 @@
 import re
 import textwrap
 from .base import Operation
+from django.utils import six
 
 
 class SeparateDatabaseAndState(Operation):
@@ -107,12 +108,17 @@ class RunPython(Operation):
     reversible = False
 
     def __init__(self, code):
+        if isinstance(code, six.string_types):
             # Trim any leading whitespace that is at the start of all code lines
             # so users can nicely indent code in migration files
             code = textwrap.dedent(code)
             # Run the code through a parser first to make sure it's at least
             # syntactically correct
             self.code = compile(code, "<string>", "exec")
+            self.is_callable = False
+        else:
+            self.code = code
+            self.is_callable = True
 
     def state_forwards(self, app_label, state):
         # RunPython objects have no state effect. To add some, combine this
@@ -124,6 +130,9 @@ class RunPython(Operation):
         # object, representing the versioned models as an AppCache.
         # We could try to override the global cache, but then people will still
         # use direct imports, so we go with a documentation approach instead.
+        if self.is_callable:
+            self.code(models=from_state.render(), schema_editor=schema_editor)
+        else:
             context = {
                 "models": from_state.render(),
                 "schema_editor": schema_editor,

tests/migrations/test_operations.py

@@ -332,6 +332,15 @@ class OperationTests(MigrationTestBase):
         # And test reversal fails
         with self.assertRaises(NotImplementedError):
             operation.database_backwards("test_runpython", None, new_state, project_state)
+        # Now test we can do it with a callable
+        def inner_method(models, schema_editor):
+            Pony = models.get_model("test_runpython", "Pony")
+            Pony.objects.create(pink=1, weight=3.55)
+            Pony.objects.create(weight=5)
+        operation = migrations.RunPython(inner_method)
+        with connection.schema_editor() as editor:
+            operation.database_forwards("test_runpython", editor, project_state, new_state)
+        self.assertEqual(project_state.render().get_model("test_runpython", "Pony").objects.count(), 4)
 
 
 class MigrateNothingRouter(object):