<s> try : <EOL> from collections import OrderedDict <EOL> except ImportError : <EOL> OrderedDict = dict <EOL> from collections import namedtuple <EOL> import re <EOL> from peewee import * <EOL> try : <EOL> from MySQLdb . constants import FIELD_TYPE <EOL> except ImportError : <EOL> try : <EOL> from pymysql . constants import FIELD_TYPE <EOL> except ImportError : <EOL> FIELD_TYPE = None <EOL> try : <EOL> from playhouse import postgres_ext <EOL> except ImportError : <EOL> postgres_ext = None <EOL> RESERVED_WORDS = set ( [ <EOL> '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT:class>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , <EOL> '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , <EOL> '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , <EOL> '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , <EOL> ] ) <EOL> class UnknownField ( object ) : <EOL> pass <EOL> class Column ( object ) : <EOL> """<STR_LIT>""" <EOL> primary_key_types = ( IntegerField , PrimaryKeyField ) <EOL> def __init__ ( self , name , field_class , raw_column_type , nullable , <EOL> primary_key = False , db_column = None , index = False , unique = False ) : <EOL> self . name = name <EOL> self . field_class = field_class <EOL> self . raw_column_type = raw_column_type <EOL> self . nullable = nullable <EOL> self . primary_key = primary_key <EOL> self . db_column = db_column <EOL> self . index = index <EOL> self . unique = unique <EOL> self . rel_model = None <EOL> self . related_name = None <EOL> self . to_field = None <EOL> def __repr__ ( self ) : <EOL> attrs = [ <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT:primary_key>' , <EOL> '<STR_LIT>' ] <EOL> keyword_args = '<STR_LIT:U+002CU+0020>' . join ( <EOL> '<STR_LIT>' % ( attr , getattr ( self , attr ) ) <EOL> for attr in attrs ) <EOL> return '<STR_LIT>' % ( self . name , keyword_args ) <EOL> def get_field_parameters ( self ) : <EOL> params = { } <EOL> if self . nullable : <EOL> params [ '<STR_LIT:null>' ] = True <EOL> if self . field_class is ForeignKeyField or self . name != self . db_column : <EOL> params [ '<STR_LIT>' ] = "<STR_LIT>" % self . db_column <EOL> if self . primary_key and self . field_class is not PrimaryKeyField : <EOL> params [ '<STR_LIT:primary_key>' ] = True <EOL> if self . is_foreign_key ( ) : <EOL> params [ '<STR_LIT>' ] = self . rel_model <EOL> if self . to_field : <EOL> params [ '<STR_LIT>' ] = "<STR_LIT>" % self . to_field <EOL> if self . related_name : <EOL> params [ '<STR_LIT:related_name>' ] = "<STR_LIT>" % self . related_name <EOL> if not self . is_primary_key ( ) : <EOL> if self . unique : <EOL> params [ '<STR_LIT>' ] = '<STR_LIT:True>' <EOL> elif self . index and not self . is_foreign_key ( ) : <EOL> params [ '<STR_LIT:index>' ] = '<STR_LIT:True>' <EOL> return params <EOL> def is_primary_key ( self ) : <EOL> return self . field_class is PrimaryKeyField or self . primary_key <EOL> def is_foreign_key ( self ) : <EOL> return self . field_class is ForeignKeyField <EOL> def is_self_referential_fk ( self ) : <EOL> return ( self . field_class is ForeignKeyField and <EOL> self . rel_model == "<STR_LIT>" ) <EOL> def set_foreign_key ( self , foreign_key , model_names , dest = None , <EOL> related_name = None ) : <EOL> self . foreign_key = foreign_key <EOL> self . field_class = ForeignKeyField <EOL> if foreign_key . dest_table == foreign_key . table : <EOL> self . rel_model = "<STR_LIT>" <EOL> else : <EOL> self . rel_model = model_names [ foreign_key . dest_table ] <EOL> self . to_field = dest and dest . name or None <EOL> self . related_name = related_name or None <EOL> def get_field ( self ) : <EOL> field_params = self . get_field_parameters ( ) <EOL> param_str = '<STR_LIT:U+002CU+0020>' . join ( '<STR_LIT>' % ( k , v ) <EOL> for k , v in sorted ( field_params . items ( ) ) ) <EOL> field = '<STR_LIT>' % ( <EOL> self . name , <EOL> self . field_class . __name__ , <EOL> param_str ) <EOL> if self . field_class is UnknownField : <EOL> field = '<STR_LIT>' % ( field , self . raw_column_type ) <EOL> return field <EOL> class Metadata ( object ) : <EOL> column_map = { } <EOL> def __init__ ( self , database ) : <EOL> self . database = database <EOL> def execute ( self , sql , * params ) : <EOL> return self . database . execute_sql ( sql , params ) <EOL> def get_columns ( self , table , schema = None ) : <EOL> metadata = OrderedDict ( <EOL> ( metadata . name , metadata ) <EOL> for metadata in self . database . get_columns ( table , schema ) ) <EOL> column_types = self . get_column_types ( table , schema ) <EOL> pk_names = self . get_primary_keys ( table , schema ) <EOL> if len ( pk_names ) == <NUM_LIT:1> : <EOL> pk = pk_names [ <NUM_LIT:0> ] <EOL> if column_types [ pk ] is IntegerField : <EOL> column_types [ pk ] = PrimaryKeyField <EOL> columns = OrderedDict ( ) <EOL> for name , column_data in metadata . items ( ) : <EOL> columns [ name ] = Column ( <EOL> name , <EOL> field_class = column_types [ name ] , <EOL> raw_column_type = column_data . data_type , <EOL> nullable = column_data . null , <EOL> primary_key = column_data . primary_key , <EOL> db_column = name ) <EOL> return columns <EOL> def get_column_types ( self , table , schema = None ) : <EOL> raise NotImplementedError <EOL> def get_foreign_keys ( self , table , schema = None ) : <EOL> return self . database . get_foreign_keys ( table , schema ) <EOL> def get_primary_keys ( self , table , schema = None ) : <EOL> return self . database . get_primary_keys ( table , schema ) <EOL> def get_indexes ( self , table , schema = None ) : <EOL> return self . database . get_indexes ( table , schema ) <EOL> class PostgresqlMetadata ( Metadata ) : <EOL> column_map = { <EOL> <NUM_LIT:16> : BooleanField , <EOL> <NUM_LIT> : BlobField , <EOL> <NUM_LIT:20> : BigIntegerField , <EOL> <NUM_LIT> : IntegerField , <EOL> <NUM_LIT> : IntegerField , <EOL> <NUM_LIT> : TextField , <EOL> <NUM_LIT> : FloatField , <EOL> <NUM_LIT> : FloatField , <EOL> <NUM_LIT> : CharField , <EOL> <NUM_LIT> : CharField , <EOL> <NUM_LIT> : DateField , <EOL> <NUM_LIT> : DateTimeField , <EOL> <NUM_LIT> : DateTimeField , <EOL> <NUM_LIT> : TimeField , <EOL> <NUM_LIT> : TimeField , <EOL> <NUM_LIT> : DecimalField , <EOL> <NUM_LIT> : TextField , <EOL> } <EOL> def __init__ ( self , database ) : <EOL> super ( PostgresqlMetadata , self ) . __init__ ( database ) <EOL> if postgres_ext is not None : <EOL> cursor = self . execute ( '<STR_LIT>' ) <EOL> results = cursor . fetchall ( ) <EOL> for oid , typname in results : <EOL> if '<STR_LIT>' in typname : <EOL> self . column_map [ oid ] = postgres_ext . JSONField <EOL> elif '<STR_LIT>' in typname : <EOL> self . column_map [ oid ] = postgres_ext . HStoreField <EOL> elif '<STR_LIT>' in typname : <EOL> self . column_map [ oid ] = postgres_ext . TSVectorField <EOL> def get_column_types ( self , table , schema ) : <EOL> column_types = { } <EOL> identifier = '<STR_LIT>' % ( schema , table ) <EOL> cursor = self . execute ( '<STR_LIT>' % identifier ) <EOL> for column_description in cursor . description : <EOL> column_types [ column_description . name ] = self . column_map . get ( <EOL> column_description . type_code , <EOL> UnknownField ) <EOL> return column_types <EOL> def get_columns ( self , table , schema = None ) : <EOL> schema = schema or '<STR_LIT>' <EOL> return super ( PostgresqlMetadata , self ) . get_columns ( table , schema ) <EOL> def get_foreign_keys ( self , table , schema = None ) : <EOL> schema = schema or '<STR_LIT>' <EOL> return super ( PostgresqlMetadata , self ) . get_foreign_keys ( table , schema ) <EOL> def get_primary_keys ( self , table , schema = None ) : <EOL> schema = schema or '<STR_LIT>' <EOL> return super ( PostgresqlMetadata , self ) . get_primary_keys ( table , schema ) <EOL> def get_indexes ( self , table , schema = None ) : <EOL> schema = schema or '<STR_LIT>' <EOL> return super ( PostgresqlMetadata , self ) . get_indexes ( table , schema ) <EOL> class MySQLMetadata ( Metadata ) : <EOL> if FIELD_TYPE is None : <EOL> column_map = { } <EOL> else : <EOL> column_map = { <EOL> FIELD_TYPE . BLOB : TextField , <EOL> FIELD_TYPE . CHAR : CharField , <EOL> FIELD_TYPE . DATE : DateField , <EOL> FIELD_TYPE . DATETIME : DateTimeField , <EOL> FIELD_TYPE . DECIMAL : DecimalField , <EOL> FIELD_TYPE . DOUBLE : FloatField , <EOL> FIELD_TYPE . FLOAT : FloatField , <EOL> FIELD_TYPE . INT24 : IntegerField , <EOL> FIELD_TYPE . LONG_BLOB : TextField , <EOL> FIELD_TYPE . LONG : IntegerField , <EOL> FIELD_TYPE . LONGLONG : BigIntegerField , <EOL> FIELD_TYPE . MEDIUM_BLOB : TextField , <EOL> FIELD_TYPE . NEWDECIMAL : DecimalField , <EOL> FIELD_TYPE . SHORT : IntegerField , <EOL> FIELD_TYPE . STRING : CharField , <EOL> FIELD_TYPE . TIMESTAMP : DateTimeField , <EOL> FIELD_TYPE . TIME : TimeField , <EOL> FIELD_TYPE . TINY_BLOB : TextField , <EOL> FIELD_TYPE . TINY : IntegerField , <EOL> FIELD_TYPE . VAR_STRING : CharField , <EOL> } <EOL> def __init__ ( self , database , ** kwargs ) : <EOL> if '<STR_LIT:password>' in kwargs : <EOL> kwargs [ '<STR_LIT>' ] = kwargs . pop ( '<STR_LIT:password>' ) <EOL> super ( MySQLMetadata , self ) . __init__ ( database , ** kwargs ) <EOL> def get_column_types ( self , table , schema = None ) : <EOL> column_types = { } <EOL> cursor = self . execute ( '<STR_LIT>' % table ) <EOL> for column_description in cursor . description : <EOL> name , type_code = column_description [ : <NUM_LIT:2> ] <EOL> column_types [ name ] = self . column_map . get ( type_code , UnknownField ) <EOL> return column_types <EOL> class SqliteMetadata ( Metadata ) : <EOL> column_map = { <EOL> '<STR_LIT>' : BigIntegerField , <EOL> '<STR_LIT>' : BlobField , <EOL> '<STR_LIT:bool>' : BooleanField , <EOL> '<STR_LIT>' : BooleanField , <EOL> '<STR_LIT>' : CharField , <EOL> '<STR_LIT:date>' : DateField , <EOL> '<STR_LIT>' : DateTimeField , <EOL> '<STR_LIT>' : DecimalField , <EOL> '<STR_LIT>' : IntegerField , <EOL> '<STR_LIT>' : IntegerField , <EOL> '<STR_LIT:int>' : IntegerField , <EOL> '<STR_LIT>' : BigIntegerField , <EOL> '<STR_LIT>' : FloatField , <EOL> '<STR_LIT>' : IntegerField , <EOL> '<STR_LIT>' : IntegerField , <EOL> '<STR_LIT>' : IntegerField , <EOL> '<STR_LIT:text>' : TextField , <EOL> '<STR_LIT:time>' : TimeField , <EOL> } <EOL> begin = '<STR_LIT>' <EOL> end = '<STR_LIT>' <EOL> re_foreign_key = ( <EOL> '<STR_LIT>' <EOL> '<STR_LIT>' <EOL> '<STR_LIT>' <EOL> '<STR_LIT>' ) . format ( begin = begin , end = end ) <EOL> re_varchar = r'<STR_LIT>' <EOL> def _map_col ( self , column_type ) : <EOL> raw_column_type = column_type . lower ( ) <EOL> if raw_column_type in self . column_map : <EOL> field_class = self . column_map [ raw_column_type ] <EOL> elif re . search ( self . re_varchar , raw_column_type ) : <EOL> field_class = CharField <EOL> else : <EOL> column_type = re . sub ( '<STR_LIT>' , '<STR_LIT>' , raw_column_type ) <EOL> field_class = self . column_map . get ( column_type , UnknownField ) <EOL> return field_class <EOL> def get_column_types ( self , table , schema = None ) : <EOL> column_types = { } <EOL> columns = self . database . get_columns ( table ) <EOL> for column in columns : <EOL> column_types [ column . name ] = self . _map_col ( column . data_type ) <EOL> return column_types <EOL> _DatabaseMetadata = namedtuple ( '<STR_LIT>' , ( <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' ) ) <EOL> class DatabaseMetadata ( _DatabaseMetadata ) : <EOL> def multi_column_indexes ( self , table ) : <EOL> accum = [ ] <EOL> for index in self . indexes [ table ] : <EOL> if len ( index . columns ) > <NUM_LIT:1> : <EOL> field_names = [ self . columns [ table ] [ column ] . name <EOL> for column in index . columns ] <EOL> accum . append ( ( field_names , index . unique ) ) <EOL> return accum <EOL> def column_indexes ( self , table ) : <EOL> accum = { } <EOL> for index in self . indexes [ table ] : <EOL> if len ( index . columns ) == <NUM_LIT:1> : <EOL> accum [ index . columns [ <NUM_LIT:0> ] ] = index . unique <EOL> return accum <EOL> class Introspector ( object ) : <EOL> pk_classes = [ PrimaryKeyField , IntegerField ] <EOL> def __init__ ( self , metadata , schema = None ) : <EOL> self . metadata = metadata <EOL> self . schema = schema <EOL> def __repr__ ( self ) : <EOL> return '<STR_LIT>' % self . metadata . database <EOL> @ classmethod <EOL> def from_database ( cls , database , schema = None ) : <EOL> if isinstance ( database , PostgresqlDatabase ) : <EOL> metadata = PostgresqlMetadata ( database ) <EOL> elif isinstance ( database , MySQLDatabase ) : <EOL> metadata = MySQLMetadata ( database ) <EOL> else : <EOL> metadata = SqliteMetadata ( database ) <EOL> return cls ( metadata , schema = schema ) <EOL> def get_database_class ( self ) : <EOL> return type ( self . metadata . database ) <EOL> def get_database_name ( self ) : <EOL> return self . metadata . database . database <EOL> def get_database_kwargs ( self ) : <EOL> return self . metadata . database . connect_kwargs <EOL> def make_model_name ( self , table ) : <EOL> model = re . sub ( '<STR_LIT>' , '<STR_LIT>' , table ) <EOL> return '<STR_LIT>' . join ( sub . title ( ) for sub in model . split ( '<STR_LIT:_>' ) ) <EOL> def make_column_name ( self , column ) : <EOL> column = re . sub ( '<STR_LIT>' , '<STR_LIT>' , column . lower ( ) ) or column . lower ( ) <EOL> if column in RESERVED_WORDS : <EOL> column += '<STR_LIT:_>' <EOL> return column <EOL> def introspect ( self , table_names = None ) : <EOL> if self . schema : <EOL> tables = self . metadata . database . get_tables ( schema = self . schema ) <EOL> else : <EOL> tables = self . metadata . database . get_tables ( ) <EOL> if table_names is not None : <EOL> tables = [ table for table in tables if table in table_names ] <EOL> columns = { } <EOL> primary_keys = { } <EOL> foreign_keys = { } <EOL> model_names = { } <EOL> indexes = { } <EOL> for table in tables : <EOL> table_indexes = self . metadata . get_indexes ( table , self . schema ) <EOL> table_columns = self . metadata . get_columns ( table , self . schema ) <EOL> try : <EOL> foreign_keys [ table ] = self . metadata . get_foreign_keys ( <EOL> table , self . schema ) <EOL> except ValueError as exc : <EOL> err ( exc . message ) <EOL> foreign_keys [ table ] = [ ] <EOL> model_names [ table ] = self . make_model_name ( table ) <EOL> lower_col_names = set ( column_name . lower ( ) <EOL> for column_name in table_columns ) <EOL> for col_name , column in table_columns . items ( ) : <EOL> new_name = self . make_column_name ( col_name ) <EOL> lower_name = col_name . lower ( ) <EOL> if lower_name . endswith ( '<STR_LIT>' ) and new_name in lower_col_names : <EOL> new_name = col_name . lower ( ) <EOL> column . name = new_name <EOL> for index in table_indexes : <EOL> if len ( index . columns ) == <NUM_LIT:1> : <EOL> column = index . columns [ <NUM_LIT:0> ] <EOL> if column in table_columns : <EOL> table_columns [ column ] . unique = index . unique <EOL> table_columns [ column ] . index = True <EOL> primary_keys [ table ] = self . metadata . get_primary_keys ( <EOL> table , self . schema ) <EOL> columns [ table ] = table_columns <EOL> indexes [ table ] = table_indexes <EOL> related_names = { } <EOL> sort_fn = lambda foreign_key : foreign_key . column <EOL> for table in tables : <EOL> models_referenced = set ( ) <EOL> for foreign_key in sorted ( foreign_keys [ table ] , key = sort_fn ) : <EOL> try : <EOL> column = columns [ table ] [ foreign_key . column ] <EOL> except KeyError : <EOL> continue <EOL> dest_table = foreign_key . dest_table <EOL> if dest_table in models_referenced : <EOL> related_names [ column ] = '<STR_LIT>' % ( <EOL> dest_table , <EOL> column . name ) <EOL> else : <EOL> models_referenced . add ( dest_table ) <EOL> for table in tables : <EOL> for foreign_key in foreign_keys [ table ] : <EOL> src = columns [ foreign_key . table ] [ foreign_key . column ] <EOL> try : <EOL> dest = columns [ foreign_key . dest_table ] [ <EOL> foreign_key . dest_column ] <EOL> except KeyError : <EOL> dest = None <EOL> src . set_foreign_key ( <EOL> foreign_key = foreign_key , <EOL> model_names = model_names , <EOL> dest = dest , <EOL> related_name = related_names . get ( src ) ) <EOL> return DatabaseMetadata ( <EOL> columns , <EOL> primary_keys , <EOL> foreign_keys , <EOL> model_names , <EOL> indexes ) <EOL> def generate_models ( self , skip_invalid = False , table_names = None ) : <EOL> database = self . introspect ( table_names = table_names ) <EOL> models = { } <EOL> class BaseModel ( Model ) : <EOL> class Meta : <EOL> database = self . metadata . database <EOL> def _create_model ( table , models ) : <EOL> for foreign_key in database . foreign_keys [ table ] : <EOL> dest = foreign_key . dest_table <EOL> if dest not in models and dest != table : <EOL> _create_model ( dest , models ) <EOL> primary_keys = [ ] <EOL> columns = database . columns [ table ] <EOL> for db_column , column in columns . items ( ) : <EOL> if column . primary_key : <EOL> primary_keys . append ( column . name ) <EOL> multi_column_indexes = database . multi_column_indexes ( table ) <EOL> column_indexes = database . column_indexes ( table ) <EOL> class Meta : <EOL> indexes = multi_column_indexes <EOL> composite_key = False <EOL> if len ( primary_keys ) == <NUM_LIT:0> : <EOL> primary_keys = columns . keys ( ) <EOL> if len ( primary_keys ) > <NUM_LIT:1> : <EOL> Meta . primary_key = CompositeKey ( * [ <EOL> field . name for col , field in columns . items ( ) <EOL> if col in primary_keys ] ) <EOL> composite_key = True <EOL> attrs = { '<STR_LIT:Meta>' : Meta } <EOL> for db_column , column in columns . items ( ) : <EOL> FieldClass = column . field_class <EOL> if FieldClass is UnknownField : <EOL> FieldClass = BareField <EOL> params = { <EOL> '<STR_LIT>' : db_column , <EOL> '<STR_LIT:null>' : column . nullable } <EOL> if column . primary_key and composite_key : <EOL> if FieldClass is PrimaryKeyField : <EOL> FieldClass = IntegerField <EOL> params [ '<STR_LIT:primary_key>' ] = False <EOL> elif column . primary_key and FieldClass is not PrimaryKeyField : <EOL> params [ '<STR_LIT:primary_key>' ] = True <EOL> if column . is_foreign_key ( ) : <EOL> if column . is_self_referential_fk ( ) : <EOL> params [ '<STR_LIT>' ] = '<STR_LIT>' <EOL> else : <EOL> dest_table = column . foreign_key . dest_table <EOL> params [ '<STR_LIT>' ] = models [ dest_table ] <EOL> if column . to_field : <EOL> params [ '<STR_LIT>' ] = column . to_field <EOL> params [ '<STR_LIT:related_name>' ] = '<STR_LIT>' % ( table , db_column ) <EOL> if db_column in column_indexes and not column . is_primary_key ( ) : <EOL> if column_indexes [ db_column ] : <EOL> params [ '<STR_LIT>' ] = True <EOL> elif not column . is_foreign_key ( ) : <EOL> params [ '<STR_LIT:index>' ] = True <EOL> attrs [ column . name ] = FieldClass ( ** params ) <EOL> try : <EOL> models [ table ] = type ( str ( table ) , ( BaseModel , ) , attrs ) <EOL> except ValueError : <EOL> if not skip_invalid : <EOL> raise <EOL> for table , model in sorted ( database . model_names . items ( ) ) : <EOL> if table not in models : <EOL> _create_model ( table , models ) <EOL> return models <EOL> def introspect ( database , schema = None ) : <EOL> introspector = Introspector . from_database ( database , schema = schema ) <EOL> return introspector . introspect ( ) </s>
