From fe9f342d8ceae96a03295e56ec743ed2e2e5fb51 Mon Sep 17 00:00:00 2001
From: Andrew Godwin <andrew@aeracode.org>
Date: Wed, 25 Sep 2013 16:10:43 +0100
Subject: [PATCH] Allow callables as the argument to RunPython

---
 django/db/migrations/operations/special.py | 31 ++++++++++++++--------
 tests/migrations/test_operations.py        |  9 +++++++
 2 files changed, 29 insertions(+), 11 deletions(-)

diff --git a/django/db/migrations/operations/special.py b/django/db/migrations/operations/special.py
index 89b13c46a0..6f7dec8c74 100644
--- a/django/db/migrations/operations/special.py
+++ b/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):
-        # 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")
+        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,11 +130,14 @@ 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.
-        context = {
-            "models": from_state.render(),
-            "schema_editor": schema_editor,
-        }
-        eval(self.code, context)
+        if self.is_callable:
+            self.code(models=from_state.render(), schema_editor=schema_editor)
+        else:
+            context = {
+                "models": from_state.render(),
+                "schema_editor": schema_editor,
+            }
+            eval(self.code, context)
 
     def database_backwards(self, app_label, schema_editor, from_state, to_state):
         raise NotImplementedError("You cannot reverse this operation")
diff --git a/tests/migrations/test_operations.py b/tests/migrations/test_operations.py
index 1f247c1a97..db79e730e2 100644
--- a/tests/migrations/test_operations.py
+++ b/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):