Fixed #22788 -- Ensured custom migration operations can be written.

This inspects the migration operation, and if it is not in the
django.db.migrations module, it adds the relevant imports to the
migration writer and uses the correct class name.

1 files changed, 10 insertions, 2 deletions

diff --git a/django/db/migrations/writer.py b/django/db/migrations/writer.py
index 88145733ec..dee867929d 100644
--- a/django/db/migrations/writer.py
+++ b/django/db/migrations/writer.py
@@ -10,7 +10,7 @@ import sys
 import types
 
 from django.apps import apps
-from django.db import models
+from django.db import models, migrations
 from django.db.migrations.loader import MigrationLoader
 from django.utils import datetime_safe, six
 from django.utils.encoding import force_text
@@ -44,7 +44,15 @@ class OperationWriter(object):
         argspec = inspect.getargspec(self.operation.__init__)
         normalized_kwargs = inspect.getcallargs(self.operation.__init__, *args, **kwargs)
 
-        self.feed('migrations.%s(' % name)
+        # See if this operation is in django.db.migrations. If it is,
+        # We can just use the fact we already have that imported,
+        # otherwise, we need to add an import for the operation class.
+        if getattr(migrations, name, None) == self.operation.__class__:
+            self.feed('migrations.%s(' % name)
+        else:
+            imports.add('import %s' % (self.operation.__class__.__module__))
+            self.feed('%s.%s(' % (self.operation.__class__.__module__, name))
+
         self.indent()
         for arg_name in argspec.args[1:]:
             arg_value = normalized_kwargs[arg_name]