Skip to content

Commit

Permalink
Fixed #24093 -- Prevented MigrationWriter to write operation kwargs t…
Browse files Browse the repository at this point in the history
…hat are not explicitly deconstructed
  • Loading branch information
MarkusH authored and carljm committed Jan 8, 2015
1 parent f487a32 commit 862ea82
Show file tree
Hide file tree
Showing 3 changed files with 175 additions and 27 deletions.
64 changes: 39 additions & 25 deletions django/db/migrations/writer.py
Original file line number Diff line number Diff line change
Expand Up @@ -46,29 +46,14 @@ def __init__(self, operation):
self.buff = []

def serialize(self):
imports = set()
name, args, kwargs = self.operation.deconstruct()
argspec = inspect.getargspec(self.operation.__init__)
normalized_kwargs = inspect.getcallargs(self.operation.__init__, *args, **kwargs)

# 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]
if (arg_name in self.operation.serialization_expand_args and
isinstance(arg_value, (list, tuple, dict))):
if isinstance(arg_value, dict):
self.feed('%s={' % arg_name)
def _write(_arg_name, _arg_value):
if (_arg_name in self.operation.serialization_expand_args and
isinstance(_arg_value, (list, tuple, dict))):
if isinstance(_arg_value, dict):
self.feed('%s={' % _arg_name)
self.indent()
for key, value in arg_value.items():
for key, value in _arg_value.items():
key_string, key_imports = MigrationWriter.serialize(key)
arg_string, arg_imports = MigrationWriter.serialize(value)
self.feed('%s: %s,' % (key_string, arg_string))
Expand All @@ -77,18 +62,47 @@ def serialize(self):
self.unindent()
self.feed('},')
else:
self.feed('%s=[' % arg_name)
self.feed('%s=[' % _arg_name)
self.indent()
for item in arg_value:
for item in _arg_value:
arg_string, arg_imports = MigrationWriter.serialize(item)
self.feed('%s,' % arg_string)
imports.update(arg_imports)
self.unindent()
self.feed('],')
else:
arg_string, arg_imports = MigrationWriter.serialize(arg_value)
self.feed('%s=%s,' % (arg_name, arg_string))
arg_string, arg_imports = MigrationWriter.serialize(_arg_value)
self.feed('%s=%s,' % (_arg_name, arg_string))
imports.update(arg_imports)

imports = set()
name, args, kwargs = self.operation.deconstruct()
argspec = inspect.getargspec(self.operation.__init__)

# 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()

# Start at one because argspec includes "self"
for i, arg in enumerate(args, 1):
arg_value = arg
arg_name = argspec.args[i]
_write(arg_name, arg_value)

i = len(args)
# Only iterate over remaining arguments
for arg_name in argspec.args[i + 1:]:
if arg_name in kwargs:
arg_value = kwargs[arg_name]
_write(arg_name, arg_value)

self.unindent()
self.feed('),')
return self.render(), imports
Expand Down
61 changes: 61 additions & 0 deletions tests/custom_migration_operations/operations.py
Original file line number Diff line number Diff line change
Expand Up @@ -31,3 +31,64 @@ def database_backwards(self, app_label, schema_editor, from_state, to_state):

class CreateModel(TestOperation):
pass


class ArgsOperation(TestOperation):
def __init__(self, arg1, arg2):
self.arg1, self.arg2 = arg1, arg2

def deconstruct(self):
return (
self.__class__.__name__,
[self.arg1, self.arg2],
{}
)


class KwargsOperation(TestOperation):
def __init__(self, kwarg1=None, kwarg2=None):
self.kwarg1, self.kwarg2 = kwarg1, kwarg2

def deconstruct(self):
kwargs = {}
if self.kwarg1 is not None:
kwargs['kwarg1'] = self.kwarg1
if self.kwarg2 is not None:
kwargs['kwarg2'] = self.kwarg2
return (
self.__class__.__name__,
[],
kwargs
)


class ArgsKwargsOperation(TestOperation):
def __init__(self, arg1, arg2, kwarg1=None, kwarg2=None):
self.arg1, self.arg2 = arg1, arg2
self.kwarg1, self.kwarg2 = kwarg1, kwarg2

def deconstruct(self):
kwargs = {}
if self.kwarg1 is not None:
kwargs['kwarg1'] = self.kwarg1
if self.kwarg2 is not None:
kwargs['kwarg2'] = self.kwarg2
return (
self.__class__.__name__,
[self.arg1, self.arg2],
kwargs,
)


class ExpandArgsOperation(TestOperation):
serialization_expand_args = ['arg']

def __init__(self, arg):
self.arg = arg

def deconstruct(self):
return (
self.__class__.__name__,
[self.arg],
{}
)
77 changes: 75 additions & 2 deletions tests/migrations/test_writer.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,8 +10,8 @@

from django.core.validators import RegexValidator, EmailValidator
from django.db import models, migrations
from django.db.migrations.writer import MigrationWriter, SettingsReference
from django.test import TestCase, ignore_warnings
from django.db.migrations.writer import MigrationWriter, OperationWriter, SettingsReference
from django.test import SimpleTestCase, TestCase, ignore_warnings
from django.conf import settings
from django.utils import datetime_safe, six
from django.utils.deconstruct import deconstructible
Expand All @@ -30,6 +30,79 @@ def upload_to(self):
thing = models.FileField(upload_to=upload_to)


class OperationWriterTests(SimpleTestCase):

def test_empty_signature(self):
operation = custom_migration_operations.operations.TestOperation()
writer = OperationWriter(operation)
writer.indentation = 0
buff, imports = writer.serialize()
self.assertEqual(imports, {'import custom_migration_operations.operations'})
self.assertEqual(
buff,
'custom_migration_operations.operations.TestOperation(\n'
'),'
)

def test_args_signature(self):
operation = custom_migration_operations.operations.ArgsOperation(1, 2)
writer = OperationWriter(operation)
writer.indentation = 0
buff, imports = writer.serialize()
self.assertEqual(imports, {'import custom_migration_operations.operations'})
self.assertEqual(
buff,
'custom_migration_operations.operations.ArgsOperation(\n'
' arg1=1,\n'
' arg2=2,\n'
'),'
)

def test_kwargs_signature(self):
operation = custom_migration_operations.operations.KwargsOperation(kwarg1=1)
writer = OperationWriter(operation)
writer.indentation = 0
buff, imports = writer.serialize()
self.assertEqual(imports, {'import custom_migration_operations.operations'})
self.assertEqual(
buff,
'custom_migration_operations.operations.KwargsOperation(\n'
' kwarg1=1,\n'
'),'
)

def test_args_kwargs_signature(self):
operation = custom_migration_operations.operations.ArgsKwargsOperation(1, 2, kwarg2=4)
writer = OperationWriter(operation)
writer.indentation = 0
buff, imports = writer.serialize()
self.assertEqual(imports, {'import custom_migration_operations.operations'})
self.assertEqual(
buff,
'custom_migration_operations.operations.ArgsKwargsOperation(\n'
' arg1=1,\n'
' arg2=2,\n'
' kwarg2=4,\n'
'),'
)

def test_expand_args_signature(self):
operation = custom_migration_operations.operations.ExpandArgsOperation([1, 2])
writer = OperationWriter(operation)
writer.indentation = 0
buff, imports = writer.serialize()
self.assertEqual(imports, {'import custom_migration_operations.operations'})
self.assertEqual(
buff,
'custom_migration_operations.operations.ExpandArgsOperation(\n'
' arg=[\n'
' 1,\n'
' 2,\n'
' ],\n'
'),'
)


class WriterTests(TestCase):
"""
Tests the migration writer (makes migration files from Migration instances)
Expand Down

0 comments on commit 862ea82

Please sign in to comment.