Added support for JSON serialization via dataclass module

Author: nathanjshaffer

sqlacodegen/codegen.py

@@ -31,6 +31,8 @@
 
 _flask_prepend = 'db.'
 
+_dataclass = False
+
 
 class _DummyInflectEngine(object):
     def singular_noun(self, noun):
@@ -312,7 +314,7 @@ def render(self):
 class ModelClass(Model):
     parent_name = 'Base'
 
-    def __init__(self, table, association_tables, inflect_engine, detect_joined):
+    def __init__(self, table, association_tables, inflect_engine, detect_joined, collector):
         super(ModelClass, self).__init__(table)
         self.name = self._tablename_to_classname(table.name, inflect_engine)
         self.children = []
@@ -321,6 +323,10 @@ def __init__(self, table, association_tables, inflect_engine, detect_joined):
         # Assign attribute names for columns
         for column in table.columns:
             self._add_attribute(column.name, column)
+            if _dataclass:
+                if column.type.python_type.__module__ != 'builtins':
+                    collector.add_literal_import(column.type.python_type.__module__, column.type.python_type.__name__)
+            
 
         # Add many-to-one relationships
         pk_column_names = set(col.name for col in table.primary_key.columns)
@@ -367,7 +373,13 @@ def add_imports(self, collector):
             child.add_imports(collector)
 
     def render(self):
+        global _dataclass        
+            
         text = 'class {0}({1}):\n'.format(self.name, self.parent_name)
+        
+        if _dataclass:
+            text = '@dataclass\n' + text
+            
         text += '    __tablename__ = {0!r}\n'.format(self.table.name)
 
         # Render constraints and indexes as __table_args__
@@ -402,6 +414,9 @@ def render(self):
         for attr, column in self.attributes.items():
             if isinstance(column, Column):
                 show_name = attr != column.name
+                if _dataclass:                    
+                    text += '    ' + attr + ' : ' + column.type.python_type.__name__  + '\n'
+                
                 text += '    {0} = {1}\n'.format(attr, _render_column(column, show_name))
 
         # Render relationships
@@ -534,7 +549,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, nocomments=False, notables=False):
+                 flask=False, ignore_cols=None, noclasses=False, nocomments=False, notables=False, dataclass=False):
         super(CodeGenerator, self).__init__()
 
         if noinflect:
@@ -552,6 +567,11 @@ def __init__(self, metadata, noindexes=False, noconstraints=False,
             _flask_prepend = ''
 
         self.nocomments = nocomments
+        
+        self.dataclass = dataclass
+        if self.dataclass:
+            global _dataclass
+            _dataclass = True
 
         # Pick association tables from the metadata into their own set, don't process them normally
         links = defaultdict(lambda: [])
@@ -610,13 +630,13 @@ def __init__(self, metadata, noindexes=False, noconstraints=False,
 
             # Only generate classes when notables is set to True
             if notables:
-                model = ModelClass(table, links[table.name], inflect_engine, not nojoined)
+                model = ModelClass(table, links[table.name], inflect_engine, not nojoined, self.collector)
                 classes[model.name] = model
             elif not table.primary_key or table.name in association_tables or noclasses:
                 # Only form model classes for tables that have a primary key and are not association tables
                 model = ModelTable(table)
             elif not noclasses:
-                model = ModelClass(table, links[table.name], inflect_engine, not nojoined)
+                model = ModelClass(table, links[table.name], inflect_engine, not nojoined, self.collector)
                 classes[model.name] = model
 
             self.models.append(model)
@@ -652,8 +672,13 @@ def __init__(self, metadata, noindexes=False, noconstraints=False,
         else:
             self.collector.add_literal_import('sqlalchemy.ext.declarative', 'declarative_base')
             self.collector.add_literal_import('sqlalchemy', 'MetaData')
+            
+            
+        if self.dataclass:
+            self.collector.add_literal_import('dataclasses', 'dataclass')
 
     def render(self, outfile=sys.stdout):
+        
         print(self.header, file=outfile)
 
         # Render the collected imports

sqlacodegen/main.py

@@ -39,6 +39,7 @@ def main():
     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")
+    parser.add_argument('--dataclass', action='store_true', help="add dataclass decorators for JSON serialization")
     args = parser.parse_args()
 
     if args.version:
@@ -58,7 +59,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.nocomments, args.notables)
+                              args.flask, ignore_cols, args.noclasses, args.nocomments, args.notables, args.dataclass)
     generator.render(outfile)