diff options
| -rw-r--r-- | django/db/migrations/operations/special.py | 26 |
1 files changed, 21 insertions, 5 deletions
diff --git a/django/db/migrations/operations/special.py b/django/db/migrations/operations/special.py index 6f7dec8c74..83cdea308e 100644 --- a/django/db/migrations/operations/special.py +++ b/django/db/migrations/operations/special.py @@ -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): - raise NotImplementedError("You cannot reverse this operation") + 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" |