Add reverse_code optional argument to RunPython
This commit is contained in:
parent
15ed75d632
commit
6d3faba2d2
|
@ -107,7 +107,8 @@ class RunPython(Operation):
|
|||
reduces_to_sql = False
|
||||
reversible = False
|
||||
|
||||
def __init__(self, code):
|
||||
def __init__(self, code, reverse_code=None):
|
||||
# Forwards 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
|
||||
|
@ -115,10 +116,16 @@ class RunPython(Operation):
|
|||
# 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
|
||||
# Reverse code
|
||||
if reverse_code is None:
|
||||
self.reverse_code = None
|
||||
elif isinstance(reverse_code, six.string_types):
|
||||
reverse_code = textwrap.dedent(reverse_code)
|
||||
self.reverse_code = compile(reverse_code, "<string>", "exec")
|
||||
else:
|
||||
self.reverse_code = reverse_code
|
||||
|
||||
def state_forwards(self, app_label, state):
|
||||
# RunPython objects have no state effect. To add some, combine this
|
||||
|
@ -130,7 +137,7 @@ 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:
|
||||
if six.callable(self.code):
|
||||
self.code(models=from_state.render(), schema_editor=schema_editor)
|
||||
else:
|
||||
context = {
|
||||
|
@ -140,7 +147,16 @@ class RunPython(Operation):
|
|||
eval(self.code, context)
|
||||
|
||||
def database_backwards(self, app_label, schema_editor, from_state, to_state):
|
||||
if self.reverse_code is None:
|
||||
raise NotImplementedError("You cannot reverse this operation")
|
||||
elif six.callable(self.reverse_code):
|
||||
self.reverse_code(models=from_state.render(), schema_editor=schema_editor)
|
||||
else:
|
||||
context = {
|
||||
"models": from_state.render(),
|
||||
"schema_editor": schema_editor,
|
||||
}
|
||||
eval(self.reverse_code, context)
|
||||
|
||||
def describe(self):
|
||||
return "Raw Python operation"
|
||||
|
|
Loading…
Reference in New Issue