Add reverse_code optional argument to RunPython

This commit is contained in:
Andrew Godwin 2013-09-25 17:20:06 +01:00
parent 15ed75d632
commit 6d3faba2d2
1 changed files with 21 additions and 5 deletions

View File

@ -107,7 +107,8 @@ class RunPython(Operation):
reduces_to_sql = False reduces_to_sql = False
reversible = False reversible = False
def __init__(self, code): def __init__(self, code, reverse_code=None):
# Forwards code
if isinstance(code, six.string_types): 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
@ -115,10 +116,16 @@ class RunPython(Operation):
# 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: else:
self.code = code 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): 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
@ -130,7 +137,7 @@ 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: if six.callable(self.code):
self.code(models=from_state.render(), schema_editor=schema_editor) self.code(models=from_state.render(), schema_editor=schema_editor)
else: else:
context = { context = {
@ -140,7 +147,16 @@ class RunPython(Operation):
eval(self.code, context) eval(self.code, context)
def database_backwards(self, app_label, schema_editor, from_state, to_state): 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") 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): def describe(self):
return "Raw Python operation" return "Raw Python operation"