Allow callables as the argument to RunPython
This commit is contained in:
parent
8a3e543f26
commit
fe9f342d8c
|
@ -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,
|
||||||
|
|
|
@ -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):
|
||||||
|
|
Loading…
Reference in New Issue