Merge pull request #33 from danerlt/master

[feat] add table column comment
ksindi · Dec 20, 2019 · ff08749 · ff08749
2 parents 90bc61a + 529c7b2
commit ff08749
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 4 deletions.
diff --git a/sqlacodegen/codegen.py b/sqlacodegen/codegen.py
@@ -160,13 +160,15 @@ def _render_column(column, show_name):
     if column.server_default:
         server_default = 'server_default=' + _flask_prepend + 'FetchedValue()'
 
+    comment = getattr(column, 'comment', None)
     return _flask_prepend + 'Column({0})'.format(', '.join(
         ([repr(column.name)] if show_name else []) +
         ([_render_column_type(column.type)] if render_coltype else []) +
         [_render_constraint(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] +
-        ([server_default] if column.server_default else [])
+        ([server_default] if column.server_default else []) +
+        (['info={!r}'.format(comment)] if comment else [])
     ))
 
 
@@ -527,7 +529,7 @@ class CodeGenerator(object):
 
     def __init__(self, metadata, noindexes=False, noconstraints=False,
                  nojoined=False, noinflect=False, nobackrefs=False,
-                 flask=False, ignore_cols=None, noclasses=False):
+                 flask=False, ignore_cols=None, noclasses=False, nocomments=False):
         super(CodeGenerator, self).__init__()
 
         if noinflect:
@@ -544,6 +546,8 @@ def __init__(self, metadata, noindexes=False, noconstraints=False,
             global _flask_prepend
             _flask_prepend = ''
 
+        self.nocomments = nocomments
+
         # Pick association tables from the metadata into their own set, don't process them normally
         links = defaultdict(lambda: [])
         association_tables = set()
@@ -654,7 +658,7 @@ def render(self, outfile=sys.stdout):
         # Render the model tables and classes
         for model in self.models:
             print('\n\n', file=outfile)
-            print(model.render().rstrip('\n').encode('utf-8'), file=outfile)
+            print(model.render().rstrip('\n'), file=outfile)
 
         if self.footer:
             print(self.footer, file=outfile)
diff --git a/sqlacodegen/main.py b/sqlacodegen/main.py
@@ -37,6 +37,7 @@ def main():
     parser.add_argument('--nobackrefs', action='store_true', help="don't include backrefs")
     parser.add_argument('--flask', action='store_true', help="use Flask-SQLAlchemy columns")
     parser.add_argument('--ignore-cols', help="Don't check foreign key constraints on specified columns (comma-separated)")
+    parser.add_argument('--nocomments', action='store_true', help="don't render column comments")
     args = parser.parse_args()
 
     if args.version:
@@ -56,7 +57,7 @@ def main():
     outfile = codecs.open(args.outfile, 'w', encoding='utf-8') if args.outfile else sys.stdout
     generator = CodeGenerator(metadata, args.noindexes, args.noconstraints,
                               args.nojoined, args.noinflect, args.nobackrefs,
-                              args.flask, ignore_cols, args.noclasses)
+                              args.flask, ignore_cols, args.noclasses, args.nocomments)
     generator.render(outfile)