Fixed missing "deferrable" and "initially" keyword arguments in ForeignKey constructs

agronholm · agronholm · commit 2f39cb1b9209 · 2013-09-01T14:59:59.000+03:00
diff --git a/sqlacodegen/codegen.py b/sqlacodegen/codegen.py
@@ -97,13 +97,25 @@ def _render_column_type(coltype):
     return text
 
 
+def _render_fk(fk):
+    text = "ForeignKey('%s.%s'" % (fk.column.table.name, fk.column.name)
+    if fk.deferrable:
+        text += ', deferrable=True'
+    if fk.initially:
+        text += ', initially=%r' % fk.initially
+    return text + ')'
+
+
 def _render_column(column, show_name):
     kwarg = []
     is_sole_pk = column.primary_key and len(column.table.primary_key) == 1
     dedicated_fks = [c for c in column.foreign_keys if len(c.constraint.columns) == 1]
     is_unique = any(isinstance(c, UniqueConstraint) and set(c.columns) == set([column])
                     for c in column.table.constraints)
     has_index = any(set(i.columns) == set([column]) for i in column.table.indexes)
+    
+    # Render the column type if there are no foreign keys on it or any of them points back to itself
+    render_coltype = not dedicated_fks or any(fk.column is column for fk in dedicated_fks)
 
     if column.key != column.name:
         kwarg.append('key')
@@ -123,8 +135,8 @@ def _render_column(column, show_name):
 
     return 'Column({0})'.format(', '.join(
         ([repr(column.name)] if show_name else []) +
-        ([_render_column_type(column.type)] if not dedicated_fks else []) +
-        [repr(x) for x in dedicated_fks] +
+        ([_render_column_type(column.type)] if render_coltype else []) +
+        [_render_fk(x) for x in dedicated_fks] +
         [repr(x) for x in column.constraints] +
         ['{0}={1}'.format(k, repr(getattr(column, k))) for k in kwarg]))
 
diff --git a/test/test_codegen.py b/test/test_codegen.py
@@ -5,7 +5,8 @@
 
 from nose.tools import eq_
 from sqlalchemy.engine import create_engine
-from sqlalchemy.schema import MetaData, Table, Column, CheckConstraint, UniqueConstraint, Index, ForeignKeyConstraint
+from sqlalchemy.schema import (MetaData, Table, Column, CheckConstraint, UniqueConstraint, Index, ForeignKey,
+                               ForeignKeyConstraint)
 from sqlalchemy.types import INTEGER, SMALLINT, VARCHAR, NUMERIC
 from sqlalchemy.dialects.postgresql.base import BIGINT, DOUBLE_PRECISION, BOOLEAN, ENUM
 from sqlalchemy.dialects.mysql.base import TINYINT
@@ -23,7 +24,7 @@ def remove_unicode_prefixes(text):
     remove_unicode_prefixes = lambda text: text
 
 
-class TestCodeGenerator(object):
+class TestModelGenerator(object):
     def setup(self):
         self.metadata = MetaData(create_engine('sqlite:///'))
 
@@ -868,4 +869,24 @@ def test_schema_table(self):
     Column('name', String),
     schema='testschema'
 )
+""")
+
+    def test_foreign_key_options(self):
+        Table(
+            'simple_items', self.metadata,
+            Column('name', VARCHAR, ForeignKey('simple_items.name', deferrable=True, initially='DEFERRED'))
+        )
+
+        eq_(self.generate_code(), """\
+# coding: utf-8
+from sqlalchemy import Column, ForeignKey, MetaData, String, Table
+
+
+metadata = MetaData()
+
+
+t_simple_items = Table(
+    'simple_items', metadata,
+    Column('name', String, ForeignKey('simple_items.name', deferrable=True, initially='DEFERRED'))
+)
 """)
