diff options
| -rw-r--r-- | migrate/versioning/genmodel.py | 37 |
1 files changed, 27 insertions, 10 deletions
diff --git a/migrate/versioning/genmodel.py b/migrate/versioning/genmodel.py index 85df627..efff67f 100644 --- a/migrate/versioning/genmodel.py +++ b/migrate/versioning/genmodel.py @@ -20,13 +20,11 @@ HEADER = """ ## File autogenerated by genmodel.py from sqlalchemy import * -meta = MetaData() """ -DECLARATIVE_HEADER = """ -## File autogenerated by genmodel.py +META_DEFINITION = "meta = MetaData()" -from sqlalchemy import * +DECLARATIVE_DEFINITION = """ from sqlalchemy.ext import declarative Base = declarative.declarative_base() @@ -125,6 +123,30 @@ class ModelGenerator(object): for name in names: yield metadata.tables.get(name) + def _genModelHeader(self, tables): + out = [] + import_index = [] + + out.append(HEADER) + + for table in tables: + for col in table.columns: + if "dialects" in col.type.__module__ and \ + col.type.__class__ not in import_index: + out.append("from " + col.type.__module__ + + " import " + col.type.__class__.__name__) + import_index.append(col.type.__class__) + + out.append("") + + if self.declarative: + out.append(DECLARATIVE_DEFINITION) + else: + out.append(META_DEFINITION) + out.append("") + + return out + def genBDefinition(self): """Generates the source code for a definition of B. @@ -134,11 +156,7 @@ class ModelGenerator(object): """ out = [] - if self.declarative: - out.append(DECLARATIVE_HEADER) - else: - out.append(HEADER) - out.append("") + out.extend(self._genModelHeader(self._get_tables(missingA=True))) for table in self._get_tables(missingA=True): out.extend(self._getTableDefn(table)) return '\n'.join(out) @@ -282,4 +300,3 @@ class ModelGenerator(object): except: trans.rollback() raise - |