<s> import peewee <EOL> from peewee import * <EOL> from flask_peewee . exceptions import ImproperlyConfigured <EOL> from flask_peewee . utils import load_class <EOL> class Database ( object ) : <EOL> def __init__ ( self , app , database = None ) : <EOL> self . app = app <EOL> self . database = database <EOL> if self . database is None : <EOL> self . load_database ( ) <EOL> self . register_handlers ( ) <EOL> self . Model = self . get_model_class ( ) <EOL> def load_database ( self ) : <EOL> self . database_config = dict ( self . app . config [ '<STR_LIT>' ] ) <EOL> try : <EOL> self . database_name = self . database_config . pop ( '<STR_LIT:name>' ) <EOL> self . database_engine = self . database_config . pop ( '<STR_LIT>' ) <EOL> except KeyError : <EOL> raise ImproperlyConfigured ( '<STR_LIT>' ) <EOL> try : <EOL> self . database_class = load_class ( self . database_engine ) <EOL> assert issubclass ( self . database_class , peewee . Database ) <EOL> except ImportError : <EOL> raise ImproperlyConfigured ( '<STR_LIT>' % self . database_engine ) <EOL> except AttributeError : <EOL> raise ImproperlyConfigured ( '<STR_LIT>' % self . database_engine ) <EOL> except AssertionError : <EOL> raise ImproperlyConfigured ( '<STR_LIT>' % self . database_engine ) <EOL> self . database = self . database_class ( self . database_name , ** self . database_config ) <EOL> def get_model_class ( self ) : <EOL> class BaseModel ( Model ) : <EOL> class Meta : <EOL> database = self . database <EOL> return BaseModel <EOL> def connect_db ( self ) : <EOL> self . database . connect ( ) <EOL> def close_db ( self , exc ) : <EOL> if not self . database . is_closed ( ) : <EOL> self . database . close ( ) <EOL> def register_handlers ( self ) : <EOL> self . app . before_request ( self . connect_db ) <EOL> self . app . teardown_request ( self . close_db ) </s>
