Skip to content

HTTPS clone URL

Subversion checkout URL

You can clone with HTTPS or Subversion.

Download ZIP
Browse files

Allow callables as the argument to RunPython

  • Loading branch information...
commit fe9f342d8ceae96a03295e56ec743ed2e2e5fb51 1 parent 8a3e543
Andrew Godwin andrewgodwin authored
31 django/db/migrations/operations/special.py
View
@@ -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 @@ def database_forwards(self, app_label, schema_editor, from_state, to_state):
# 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")
9 tests/migrations/test_operations.py
View
@@ -332,6 +332,15 @@ def test_run_python(self):
# 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):
Please sign in to comment.
Something went wrong with that request. Please try again.