Commit 6d3faba2 authored by Andrew Godwin's avatar Andrew Godwin
Browse files

Add reverse_code optional argument to RunPython

parent 15ed75d6
Loading
Loading
Loading
Loading
+21 −5
Original line number Diff line number Diff line
@@ -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"