Anonymous avatar Anonymous committed 4fee724

startmigration now outputs new Field based migration files.
- removed the --all option, this is now replaced with --initial to be more descriptive over its intended use
- added a --model option which allows a user to specify one or more models to include in the migration

TODO/Known Bugs:
- Many-To-Many relationships aren't output correctly yet
- Foreign Keys that use strings to identify a model they point to, whether it be 'self', or 'modelName' or 'appName.modelName' aren't interpreted correctly yet.

Comments (0)

Files changed (1)

management/commands/startmigration.py

 from django.core.management.base import BaseCommand
 from django.core.management.color import no_style
 from django.db import models
-from django.db.models.loading import get_models
 from optparse import make_option
 from south import migration
 import sys
 import os
+import re
 import string
 import random
+import inspect
+import parser
 
 class Command(BaseCommand):
     option_list = BaseCommand.option_list + (
-        make_option('--all', action='store_true', dest='all_models', default=False,
-            help='Tells South to create migrations for all models in the given app.'),
+        make_option('--model', action='append', dest='model_list', type='string',
+            help='Generate a Create Table migration for the specified model.  Add multiple models to this migration with subsequent --model parameters.'),
+        make_option('--initial', action='store_true', dest='initial', default=False,
+            help='Generate the initial schema for the app.'),
     )
     help = "Creates a new template migration for the given app"
-
-    def handle(self, app=None, name="", model=None, all_models=False, **options):
+    
+    def handle(self, app=None, name="", model_list=[], initial=False, **options):
+        # make sure --model and --all aren't both specified
+        if initial and model_list:
+            print "You cannot use --initial and other options together"
+            return
+            
+        # specify the default name 'initial' if a name wasn't specified and we're
+        # doing a migration for an entire app
+        if not name and initial:
+            name = 'initial'
+            
+        # if not name, there's an error
+        if not name:
+            print "You must name this migration"
+            return
+        
         if not app:
             print "Please provide an app in which to create the migration."
             return
+            
         # See if the app exists
+        app_models_module = models.get_app(app)
+        if not app_models_module:
+            print "App '%s' doesn't seem to exist, isn't in INSTALLED_APPS, or has no models." % app
+            return
+            
+        # Determine what models should be included in this migration.
+        models_to_migrate = []
+        if initial:
+            models_to_migrate = models.get_models(app_models_module)
+            if not models_to_migrate:
+                print "No models found in app '%s'" % (app)
+                return
+        else:
+            for model_name in model_list:
+                model = models.get_model(app, model_name)
+                if not model:
+                    print "Couldn't find model '%s' in app '%s'" % (model_name, app)
+                    return
+                    
+                models_to_migrate.append(model)
+                
+        # Make the migrations directory if it's not there
+        app_module_path = app_models_module.__name__.split('.')[0:-1]
         try:
-            app_module = __import__(app, {}, {}, ['migrations','models'])
+            app_module = __import__('.'.join(app_module_path), {}, {}, [''])
         except ImportError:
-            print "App '%s' doesn't seem to exist." % app
+            print "Couldn't find path to App '%s'." % app
             return
-        # If they asked for all models, they better not have given a name
-        if all_models:
-            if name:
-                print "You cannot ask for all models to be autogenerated, and give a migration name."
-                return
-            else:
-                # OK, go ahead and scan the models then
-                models = get_models(app_module.models)
-                for model in models:
-                    self.handle(app, model.__name__.lower(), model.__name__, **options)
-                return
-        # If there's a model, open it and have a poke
-        if model:
-            try:
-                model = getattr(app_module.models, model)
-            except AttributeError:
-                print "The specified model '%s' doesn't seem to exist." % model
-        # Make the migrations directory if it's not there
+            
         migrations_dir = os.path.join(
             os.path.dirname(app_module.__file__),
             "migrations",
             name,
         )
         # If there's a model, make the migration skeleton, else leave it bare
-        if model:
-            table_name = model._meta.db_table
-            fields = []
-            m2ms = []
-            for f in model._meta.local_fields:
-                # We use a list of tuples to get nice ordering
-                type, type_param = genericify_type(f.db_type())
-                field = [
-                    ("name", f.column),
-                    ("type", type),
-                    ("type_param", type_param),
-                    ("null", f.null),
-                ]
-                if f.primary_key:
-                    field.append(('primary', True))
-                if f.unique:
-                    field.append(('unique', True))
-                if f.rel:
-                    field.append(('related_to', (
-                        f.rel.to._meta.db_table,
-                        f.rel.to._meta.get_field(f.rel.field_name).column,
-                    )))
-                fields.append(field)
-            forwards = '''db.create_table("%s", [
+        forwards, backwards = '', ''
+        if models_to_migrate:
+            for model in models_to_migrate:
+                table_name = model._meta.db_table
+                fields = []
+                for f in model._meta.local_fields:
+                    # We use a list of tuples to get nice ordering
+                    field_definition = generate_field_definition(model, f)
+                    fields.append((f.name, field_definition))
+                forwards += '''
+        # Model '%s'
+        db.create_table("%s", (
             %s
-        ])''' % (
-                table_name,
-                ",\n            ".join([
-                    "{%s}" % ", ".join("%r: %r" % (x, y) for x, y in f)
-                    for f in fields
-                ]),
-            )
-            backwards = '''db.delete_table("%s")''' % table_name
-            # Now go through local M2Ms and add extra stuff for them
-            for m in model._meta.local_many_to_many:
-                forwards += '''\n        # M2M field '%s'
-        db.create_table("%s", [
-            {"name": "id", "type": "serial", "null": False, "unique": True},
-            {"name": "%s", "type": "integer", "null": False, "related_to": ("%s", "%s")},
-            {"name": "%s", "type": "integer", "null": False, "related_to": ("%s", "%s")},
-        ]) ''' % (
-                    m.name,
-                    m.m2m_db_table(),
-                    m.m2m_column_name(),
+        ))''' % (
+                    model._meta.object_name,
                     table_name,
-                    "id",
-                    m.m2m_reverse_name(),
-                    m.rel.to._meta.db_table,
-                    "id",
+                    ",\n            ".join(["('%s', %s)" % (f[0], f[1]) for f in fields]),
                 )
+
                 backwards += '''
-        db.delete_table("%s")''' % m.m2m_db_table()
+        db.delete_table("%s")''' % table_name
+                # Now go through local M2Ms and add extra stuff for them
+        #         for m in model._meta.local_many_to_many:
+        #             forwards += '''
+        # # M2M field '%s'
+        # db.create_table("%s", [
+        #     {"name": "id", "type": "serial", "null": False, "unique": True},
+        #     {"name": "%s", "type": "integer", "null": False, "related_to": ("%s", "%s")},
+        #     {"name": "%s", "type": "integer", "null": False, "related_to": ("%s", "%s")},
+        # ]) ''' % (
+        #                 m.name,
+        #                 m.m2m_db_table(),
+        #                 m.m2m_column_name(),
+        #                 table_name,
+        #                 "id",
+        #                 m.m2m_reverse_name(),
+        #                 m.rel.to._meta.db_table,
+        #                 "id",
+        #         )
+        #         
+        #             backwards += '''
+        # db.delete_table("%s")''' % m.m2m_db_table()
                 
             forwards += '''
         
-        db.send_create_signal('%s', ['%s'])''' % (model._meta.app_label, model._meta.object_name)
+        db.send_create_signal('%s', ['%s'])''' % (
+                app, 
+                "','".join(model._meta.object_name for model in models_to_migrate)
+                )
         
         else:
             forwards = '"Write your forwards migration here"'
     
     def backwards(self):
         %s
-""" % (app, forwards, backwards))
+""" % ('.'.join(app_module_path), forwards, backwards))
         fp.close()
         print "Created %s." % new_filename
 
 
-def genericify_type(typestr):
-    if "(" not in typestr:
-        type = typestr
-        param = None
-    else:
-        type, param = typestr.split("(")
-        param = param[:-1]
-    # Make sure it doesn't need to be mapped back to a more generic type
-    type = {
-        "varchar": "string",
-        "timestamp with time zone": "datetime",
-        "double precision": "float",
-    }.get(type, type)
-    return type, param
+def generate_field_definition(model, field):
+    """
+    Inspects the source code of 'model' to find the code used to generate 'field'
+    """
+    def test_field(field_definition):
+        try:
+            parser.suite(field_definition)
+            return True
+        except SyntaxError:
+            return False
+    
+    field_pieces = []
+    found_field = False
+    source = inspect.getsourcelines(model)
+    if not source:
+        raise Exception("Could not find source to model: '%s'" % (model.__name__))
+        
+    # look for a line starting with the field name
+    start_field_re = re.compile(r'\s*%s\s*=\s*(.*)' % field.name)
+    for line in source[0]:
+        # if the field was found during a previous iteration, 
+        # we're here because the field spans across multiple lines
+        # append the current line and try again
+        if found_field:
+            field_pieces.append(line.strip())
+            if test_field(' '.join(field_pieces)):
+                return ' '.join(field_pieces)
+            continue
+        
+        match = start_field_re.match(line)
+        if match:
+            found_field = True
+            field_pieces.append(match.groups()[0])
+            if test_field(' '.join(field_pieces)):
+                return ' '.join(field_pieces)
+                
+    # the 'id' field never gets defined, so return what django does by default
+    # django.db.models.options::_prepare
+    if field.name == 'id' and isinstance(field, models.AutoField):
+        return "models.AutoField(verbose_name='ID', primary_key=True, auto_created=True)"
+            
+    raise Exception("Couldn't find field definition for field: '%s' on model: '%s'" % (
+        field.name, model))
+
Tip: Filter by directory path e.g. /media app.js to search for public/media/app.js.
Tip: Use camelCasing e.g. ProjME to search for ProjectModifiedEvent.java.
Tip: Filter by extension type e.g. /repo .js to search for all .js files in the /repo directory.
Tip: Separate your search with spaces e.g. /ssh pom.xml to search for src/ssh/pom.xml.
Tip: Use ↑ and ↓ arrow keys to navigate and return to view the file.
Tip: You can also navigate files with Ctrl+j (next) and Ctrl+k (previous) and view the file with Ctrl+o.
Tip: You can also navigate files with Alt+j (next) and Alt+k (previous) and view the file with Alt+o.