Allow callables as the argument to RunPython

import re

import textwrap

from .base import Operation

from django.utils import six

class SeparateDatabaseAndState(Operation):

    reversible = False

    def __init__(self, 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

            code = textwrap.dedent(code)

            # 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

    def state_forwards(self, app_label, state):

        # RunPython objects have no state effect. To add some, combine this

        # 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:

            self.code(models=from_state.render(), schema_editor=schema_editor)

        else:

            context = {

                "models": from_state.render(),

                "schema_editor": schema_editor,

        # And test reversal fails

        with self.assertRaises(NotImplementedError):

            operation.database_backwards("test_runpython", None, new_state, project_state)

        # Now test we can do it with a callable

        def inner_method(models, schema_editor):

            Pony = models.get_model("test_runpython", "Pony")

            Pony.objects.create(pink=1, weight=3.55)

            Pony.objects.create(weight=5)

        operation = migrations.RunPython(inner_method)

        with connection.schema_editor() as editor:

            operation.database_forwards("test_runpython", editor, project_state, new_state)

        self.assertEqual(project_state.render().get_model("test_runpython", "Pony").objects.count(), 4)

class MigrateNothingRouter(object):