Fixed missing table argument for indicating the table's schema (fixes agronholm#2)

agronholm · agronholm · commit 3ccdae43a491 · 2013-09-01T01:24:14.000+03:00
diff --git a/sqlacodegen/codegen.py b/sqlacodegen/codegen.py
@@ -166,6 +166,7 @@ class Model(object):
     def __init__(self, table):
         super(Model, self).__init__()
         self.table = table
+        self.schema = table.schema
 
         # Adapt column types to the most reasonable generic types (ie. VARCHAR -> String)
         for column in table.columns:
@@ -224,6 +225,9 @@ def render(self):
             if len(index.columns) > 1:
                 text += '    {0},\n'.format(_render_index(index))
 
+        if self.schema:
+            text += "    schema='{0}',\n".format(self.schema)
+
         return text.rstrip('\n,') + '\n)'
 
 
@@ -293,7 +297,18 @@ def render(self):
         for index in self.table.indexes:
             if len(index.columns) > 1:
                 table_args.append(_render_index(index))
-        if table_args:
+
+        table_kwargs = {}
+        if self.schema:
+            table_kwargs['schema'] = self.schema
+
+        kwargs_items = ', '.join('{0!r}: {1!r}'.format(key, table_kwargs[key]) for key in table_kwargs)
+        kwargs_items = '{{{0}}}'.format(kwargs_items) if kwargs_items else None
+        if table_kwargs and not table_args:
+            text += '    __table_args__ = {0}\n'.format(kwargs_items)
+        elif table_args:
+            if kwargs_items:
+                table_args.append(kwargs_items)
             if len(table_args) == 1:
                 table_args[0] += ','
             text += '    __table_args__ = (\n        {0}\n    )\n'.format(',\n        '.join(table_args))
diff --git a/test/test_codegen.py b/test/test_codegen.py
@@ -26,7 +26,7 @@ def remove_unicode_prefixes(text):
 class TestCodeGenerator(object):
     def setup(self):
         self.metadata = MetaData(create_engine('sqlite:///'))
-    
+
     def generate_code(self, **kwargs):
         codegen = CodeGenerator(self.metadata, **kwargs)
         sio = StringIO()
@@ -793,3 +793,79 @@ class SimpleItems(Base):
 
     id = Column(Integer, primary_key=True)
 """)
+
+    def test_table_kwargs(self):
+        Table(
+            'simple_items', self.metadata,
+            Column('id', INTEGER, primary_key=True),
+            schema='testschema'
+        )
+
+        eq_(self.generate_code(), """\
+# coding: utf-8
+from sqlalchemy import Column, Integer
+from sqlalchemy.ext.declarative import declarative_base
+
+
+Base = declarative_base()
+metadata = Base.metadata
+
+
+class SimpleItem(Base):
+    __tablename__ = 'simple_items'
+    __table_args__ = {'schema': 'testschema'}
+
+    id = Column(Integer, primary_key=True)
+""")
+
+    def test_table_args_kwargs(self):
+        Table(
+            'simple_items', self.metadata,
+            Column('id', INTEGER, primary_key=True),
+            Column('name', VARCHAR),
+            Index('testidx', 'id', 'name'),
+            schema='testschema'
+        )
+
+        eq_(self.generate_code(), """\
+# coding: utf-8
+from sqlalchemy import Column, Index, Integer, String
+from sqlalchemy.ext.declarative import declarative_base
+
+
+Base = declarative_base()
+metadata = Base.metadata
+
+
+class SimpleItem(Base):
+    __tablename__ = 'simple_items'
+    __table_args__ = (
+        Index('testidx', 'id', 'name'),
+        {'schema': 'testschema'}
+    )
+
+    id = Column(Integer, primary_key=True)
+    name = Column(String)
+""")
+
+    def test_schema_table(self):
+        Table(
+            'simple_items', self.metadata,
+            Column('name', VARCHAR),
+            schema='testschema'
+        )
+
+        eq_(self.generate_code(), """\
+# coding: utf-8
+from sqlalchemy import Column, MetaData, String, Table
+
+
+metadata = MetaData()
+
+
+t_simple_items = Table(
+    'simple_items', metadata,
+    Column('name', String),
+    schema='testschema'
+)
+""")
