Login

Command to dump data as a python script

Author:
willhardy
Posted:
June 20, 2008
Language:
Python
Version:
.96
Score:
9 (after 9 ratings)

This creates a fixture in the form of a python script.

Handles:

  1. ForeignKey and ManyToManyFields (using python variables, not IDs)
  2. Self-referencing ForeignKey (and M2M) fields
  3. Sub-classed models
  4. ContentType fields
  5. Recursive references
  6. AutoFields are excluded
  7. Parent models are only included when no other child model links to it

There are a few benefits to this:

  1. edit script to create 1,000s of generated entries using for loops, python modules etc.
  2. little drama with model evolution: foreign keys handled naturally without IDs, new and removed columns are ignored

The runscript command by poelzi, complements this command very nicely! e.g.

$ ./manage.py dumpscript appname > scripts/testdata.py
$ ./manage.py reset appname
$ ./manage.py runscript testdata
  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
#!/usr/bin/env python
# -*- coding: UTF-8 -*-
"""
      Title: Dumpscript management command
    Project: Hardytools (queryset-refactor version)
     Author: Will Hardy (http://willhardy.com.au)
       Date: June 2008
      Usage: python manage.py dumpscript appname > scripts/scriptname.py
  $Revision: 217 $

Description: 
    Generates a Python script that will repopulate the database using objects.
    The advantage of this approach is that it is easy to understand, and more
    flexible than directly populating the database, or using XML.

    * It also allows for new defaults to take effect and only transfers what is
      needed.
    * If a new database schema has a NEW ATTRIBUTE, it is simply not
      populated (using a default value will make the transition smooth :)
    * If a new database schema REMOVES AN ATTRIBUTE, it is simply ignored
      and the data moves across safely (I'm assuming we don't want this
      attribute anymore.
    * Problems may only occur if there is a new model and is now a required
      ForeignKey for an existing model. But this is easy to fix by editing the
      populate script :)

Improvements:
    See TODOs and FIXMEs scattered throughout :-)

"""

import sys
from django.db import models
from django.core.exceptions import ObjectDoesNotExist
from django.core.management.base import BaseCommand
from django.utils.encoding import smart_unicode
from django.contrib.contenttypes.models import ContentType

class Command(BaseCommand):
    help = 'Dumps the data as a customised python script.'
    args = '[appname ...]'

    def handle(self, *app_labels, **options):

        # Get the models we want to export
        models = get_models(app_labels)

        # A dictionary is created to keep track of all the processed objects,
        # so that foreign key references can be made using python variable names.
        # This variable "context" will be passed around like the town bicycle.
        context = {}

        # Create a dumpscript object and let it format itself as a string
        print Script(models=models, context=context)


def get_models(app_labels):
    """ Gets a list of models for the given app labels, with some exceptions. 
    """

    from django.db.models import get_app, get_apps
    from django.db.models import get_models as get_all_models

    # These models are not to be output, e.g. because they can be generated automatically
    # TODO: This should also specify the app, in case model
    EXCLUDED_MODELS = ('ContentType', )

    # Get all relevant apps
    if not app_labels:
        app_list = get_apps()
    else:
        app_list = [ get_app(app_label) for app_label in app_labels ]

    # Get a list of all the relevant models
    models = []
    for app in app_list:
        # Get all models for each app, except any excluded ones
        models += [ m for m in get_all_models(app) if m.__name__ not in EXCLUDED_MODELS ]

    return models




class Code(object):
    """ A snippet of python script. 
        This keeps track of import statements and can be output to a string.
        In the future, other features such as custom indentation might be included
        in this class.
    """

    def __init__(self):
        self.imports = {}
        self.indent = -1 

    def __str__(self):
        """ Returns a string representation of this script. 
        """
        if self.imports:
            sys.stderr.write(repr(self.import_lines))
            return flatten_blocks([""] + self.import_lines + [""] + self.lines, num_indents=self.indent)
        else:
            return flatten_blocks(self.lines, num_indents=self.indent)

    def get_import_lines(self):
        """ Takes the stored imports and converts them to lines
        """
        if self.imports:
            return [ "from %s import %s" % (value, key) for key, value in self.imports.items() ]
        else:
            return []
    import_lines = property(get_import_lines)


class ModelCode(Code):
    " Produces a python script that can recreate data for a given model class. "

    def __init__(self, model, context={}):
        self.model = model
        self.context = context
        self.instances = []
        self.indent = 0

    def get_imports(self):
        """ Returns a dictionary of import statements, with the variable being
            defined as the key. 
        """
        return { self.model.__name__: smart_unicode(self.model.__module__) }
    imports = property(get_imports)

    def get_lines(self):
        """ Returns a list of lists or strings, representing the code body. 
            Each list is a block, each string is a statement.
        """
        code = []

        for counter, item in enumerate(self.model.objects.all()):
            instance = InstanceCode(instance=item, id=counter+1, context=self.context)
            self.instances.append(instance)
            if instance.waiting_list:
                code += instance.lines
 
        # After each instance has been processed, try again.
        # This allows self referencing fields to work.
        for instance in self.instances:
            if instance.waiting_list:
                code += instance.lines

        return code

    lines = property(get_lines)


class InstanceCode(Code):
    " Produces a python script that can recreate data for a given model instance. "

    def __init__(self, instance, id, context={}):
        """ We need the instance in question and an id """

        self.instance = instance
        self.model = self.instance.__class__
        self.context = context
        self.variable_name = "%s_%s" % (self.instance._meta.db_table, id)
        self.skip_me = None
        self.instantiated = False

        self.indent  = 0 
        self.imports = {}

        self.waiting_list = list(self.model._meta.fields)

        self.many_to_many_waiting_list = {} 
        for field in self.model._meta.many_to_many:
            self.many_to_many_waiting_list[field] = list(getattr(self.instance, field.name).all())

    def get_lines(self):
        """ Returns a list of lists or strings, representing the code body. 
            Each list is a block, each string is a statement.
        """
        code_lines = []

        # Don't return anything if this is an instance that should be skipped
        if self.skip():
            return []

        # Initialise our new object
        # e.g. model_name_35 = Model()
        code_lines += self.instantiate()

        # Add each field
        # e.g. model_name_35.field_one = 1034.91
        #      model_name_35.field_two = "text"
        code_lines += self.get_waiting_list()

        # Print the save command for our new object
        # e.g. model_name_35.save()
        if code_lines:
            code_lines.append("%s.save()\n" % (self.variable_name))

        code_lines += self.get_many_to_many_lines()

        return code_lines
    lines = property(get_lines)

    def skip(self):
        """ Determine whether or not this object should be skipped.
            If this model is a parent of a single subclassed instance, skip it.
            The subclassed instance will create this parent instance for us.

            TODO: Allow the user to force its creation?
        """

        if self.skip_me is not None:
            return self.skip_me

        try:
            # Django trunk since r7722 uses CollectedObjects instead of dict
            from django.db.models.query import CollectedObjects
            sub_objects = CollectedObjects()
        except ImportError:
            # previous versions don't have CollectedObjects
            sub_objects = {}
        self.instance._collect_sub_objects(sub_objects)
        if reduce(lambda x, y: x+y, [self.model in so._meta.parents for so in sub_objects.keys()]) == 1:
            pk_name = self.instance._meta.pk.name
            key = '%s_%s' % (self.model.__name__, getattr(self.instance, pk_name))
            self.context[key] = None
            self.skip_me = True
        else:
            self.skip_me = False

        return self.skip_me

    def instantiate(self):
        " Write lines for instantiation "
        # e.g. model_name_35 = Model()
        code_lines = []

        if not self.instantiated:
            code_lines.append("%s = %s()" % (self.variable_name, self.model.__name__))
            self.instantiated = True

            # Store our variable name for future foreign key references
            pk_name = self.instance._meta.pk.name
            key = '%s_%s' % (self.model.__name__, getattr(self.instance, pk_name))
            self.context[key] = self.variable_name

        return code_lines


    def get_waiting_list(self):
        " Add lines for any waiting fields that can be completed now. "

        code_lines = []

        # Process normal fields
        for field in list(self.waiting_list):
            try:
                # Find the value, add the line, remove from waiting list and move on
                value = get_attribute_value(self.instance, field, self.context)
                code_lines.append('%s.%s = %s' % (self.variable_name, field.name, value))
                self.waiting_list.remove(field)
            except SkipValue, e:
                # Remove from the waiting list and move on
                self.waiting_list.remove(field)
                continue
            except DoLater, e:
                # Move on, maybe next time
                continue


        return code_lines

    def get_many_to_many_lines(self):
        """ Generates lines that define many to many relations for this instance. """

        lines = []

        for field, rel_items in self.many_to_many_waiting_list.items():
            for rel_item in list(rel_items):
                try:
                    pk_name = rel_item._meta.pk.name
                    key = '%s_%s' % (rel_item.__class__.__name__, getattr(rel_item, pk_name))
                    value = "%s" % self.context[key]
                    lines.append('%s.%s.add(%s)' % (self.variable_name, field.name, value))
                    self.many_to_many_waiting_list[field].remove(rel_item)
                except KeyError:
                    pass

        if lines:
            lines.append("")

        return lines


class Script(Code):
    " Produces a complete python script that can recreate data for the given apps. "

    def __init__(self, models, context={}):
        self.models = models
        self.context = context

        self.indent = -1 
        self.imports = {}

    def get_lines(self):
        """ Returns a list of lists or strings, representing the code body. 
            Each list is a block, each string is a statement.
        """
        code = [ self.FILE_HEADER.strip() ]

        # Queue and process the required models
        for model_class in queue_models(self.models, context=self.context):
            sys.stderr.write('Processing model: %s\n' % model_class.model.__name__)
            code.append(model_class.import_lines)
            code.append("")
            code.append(model_class.lines)

        # Process left over foreign keys from cyclic models
        for model in self.models:
            sys.stderr.write('Re-processing model: %s\n' % model.model.__name__)
            for instance in model.instances:
                if instance.waiting_list:
                    code.append(instance.lines)

        return code

    lines = property(get_lines)

    # A user-friendly file header
    FILE_HEADER = """

#!/usr/bin/env python
# -*- coding: utf-8 -*-

# This file has been automatically generated, changes may be lost if you
# go and generate it again. It was generated with the following command:
# %s

import datetime
from decimal import Decimal
from django.contrib.contenttypes.models import ContentType

def run():

""" % " ".join(sys.argv)



# HELPER FUNCTIONS
#-------------------------------------------------------------------------------

def flatten_blocks(lines, num_indents=-1):
    """ Takes a list (block) or string (statement) and flattens it into a string
        with indentation. 
    """

    # The standard indent is four spaces
    INDENTATION = " " * 4

    if not lines:
        return ""

    # If this is a string, add the indentation and finish here
    if isinstance(lines, basestring):
        return INDENTATION * num_indents + lines

    # If this is not a string, join the lines and recurse
    return "\n".join([ flatten_blocks(line, num_indents+1) for line in lines ])




def get_attribute_value(item, field, context):
    """ Gets a string version of the given attribute's value, like repr() might. """

    # Find the value of the field, catching any database issues
    try:
        value = getattr(item, field.name)
    except ObjectDoesNotExist:
        raise SkipValue('Could not find object for %s.%s, ignoring.\n' % (item.__class__.__name__, field.name))

    # AutoField: We don't include the auto fields, they'll be automatically recreated
    if isinstance(field, models.AutoField):
        raise SkipValue()#"AutoField %s.%s" % (item.__class__, field.name))

    # ForeignKey fields, link directly using our stored python variable name
    elif isinstance(field, models.ForeignKey) and value is not None:

        # Special case for contenttype foreign keys: no need to output any
        # content types in this script, as they can be generated again 
        # automatically.
        # NB: Not sure if "is" will always work
        if field.rel.to is ContentType:
            return 'ContentType.objects.get(app_label="%s", model="%s")' % (value.app_label, value.model)

        # Generate an identifier (key) for this foreign object
        pk_name = value._meta.pk.name
        key = '%s_%s' % (value.__class__.__name__, getattr(value, pk_name))

        if key in context:
            variable_name = context[key]
            # If the context value is set to None, this should be skipped.
            # This identifies models that have been skipped (inheritance)
            if variable_name is None:
                raise SkipValue()
            # Return the variable name listed in the context 
            return "%s" % variable_name
        else:
            raise DoLater('(FK) %s.%s\n' % (item.__class__.__name__, field.name))


    # A normal field (e.g. a python built-in)
    else:
        return repr(value)

def queue_models(models, context):
    """ Works an an appropriate ordering for the models.
        This isn't essential, but makes the script look nicer because 
        more instances can be defined on their first try.
    """

    # Max number of cycles allowed before we call it an infinite loop.
    MAX_CYCLES = 5

    model_queue = []
    number_remaining_models = len(models)
    allowed_cycles = MAX_CYCLES

    while number_remaining_models > 0:
        previous_number_remaining_models = number_remaining_models

        model = models.pop(0)
        
        # If the model is ready to be processed, add it to the list
        if check_dependencies(model, model_queue):
            model_class = ModelCode(model=model, context=context)
            model_queue.append(model_class)

        # Otherwise put the model back at the end of the list
        else:
            models.append(model)

        # Check for infinite loops. 
        # This means there is a cyclic foreign key structure
        # That cannot be resolved by re-ordering
        number_remaining_models = len(models)
        if number_remaining_models == previous_number_remaining_models:
            allowed_cycles -= 1
            if allowed_cycles <= 0:
                # Add the remaining models, but do not remove them from the model list
                missing_models = [ ModelCode(model=m, context=context) for m in models ]
                model_queue += missing_models
                # Replace the models with the model class objects 
                # (sure, this is a little bit of hackery)
                models[:] = missing_models
                break
        else:
            allowed_cycles = MAX_CYCLES

    return model_queue


def check_dependencies(model, model_queue):
    " Check that all the depenedencies for this model are already in the queue. "

    # A list of allowed links: existing fields, itself and the special case ContentType
    allowed_links = [ m.model.__name__ for m in model_queue ] + [model.__name__, 'ContentType']

    # For each ForeignKey or ManyToMany field, check that a link is possible
    for field in model._meta.fields + model._meta.many_to_many:
        if field.rel and field.rel.to.__name__ not in allowed_links:
            return False

    return True



# EXCEPTIONS
#-------------------------------------------------------------------------------

class SkipValue(Exception):
    """ Value could not be parsed or should simply be skipped. """

class DoLater(Exception):
    """ Value could not be parsed or should simply be skipped. """

More like this

  1. Template tag - list punctuation for a list of items by shapiromatron 2 months ago
  2. JSONRequestMiddleware adds a .json() method to your HttpRequests by cdcarter 2 months, 1 week ago
  3. Serializer factory with Django Rest Framework by julio 9 months, 1 week ago
  4. Image compression before saving the new model / work with JPG, PNG by Schleidens 9 months, 4 weeks ago
  5. Help text hyperlinks by sa2812 10 months, 3 weeks ago

Comments

jokull (on June 24, 2008):

+1 for ./manage.py dumpdata --format python

#

trbs (on June 24, 2008):

dumpscript is now part of django-command-extensions :)

#

gorans (on June 25, 2008):

Nice one Will!

#

buriy (on June 26, 2008):

nice one. how about object and model cherry-picking? ;)

#

willhardy (on July 1, 2008):

object and model cherry-picking?

./manage.py dumpscript | grep "appname_modelname_"

But please no bug reports on this one :-P

#

aj (on July 17, 2008):

Great script Will - thanks.

I found (and fixed) a bug against the latest trunk and have added model include and exclude options on the command line. Would you like a copy so you can choose whether to update your snippet and perhaps point out where what I've done screws everything up? ;) If so I'm on aj 'at' cubbyhole 'dot' net.

AJ

#

tjurewicz (on July 30, 2008):

I'm relatively new to Django and am trying to get this script to work. Can somebody point me in a direction to running these types of scripts?

Thanks, Trent

#

johnboxall (on February 27, 2009):

Awesome but don't try this one with more than a handful of records unless you've got something beefy!

#

Please login first before commenting.