Add reverse_code optional argument to RunPython

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"