Allow callables as the argument to RunPython

This commit is contained in:
Andrew Godwin 2013-09-25 16:10:43 +01:00
parent 8a3e543f26
commit fe9f342d8c
2 changed files with 29 additions and 11 deletions

View File

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

View File

@ -332,6 +332,15 @@ class OperationTests(MigrationTestBase):
# And test reversal fails # And test reversal fails
with self.assertRaises(NotImplementedError): with self.assertRaises(NotImplementedError):
operation.database_backwards("test_runpython", None, new_state, project_state) 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): class MigrateNothingRouter(object):