Changed methods and added database class

src/database/database.py

@@ -1,8 +1,16 @@
 import sqlite3
 from sqlite3 import Error
 
+class database:
+    """Main database class that controls all the functionality for input /
+        out of the database."""
 
-def create_connection(db_file):
+
+    def __init__(self):
+        self.conn = None
+
+
+    def create_connection(self, db_file):
         """ create a database connection to the SQLite database
             specified by db_file
         :param db_file: database file
@@ -15,22 +23,22 @@ def create_connection(db_file):
         except Error as e:
             print(e)
 
-    return conn
+        self.conn = conn
 
 
-def create_table(conn, create_table_sql):
+    def create_table(self, create_table_sql):
         """ create a table from the create_table_sql statement
         :param conn: Connection object
         :param create_table_sql: a CREATE TABLE statement
         :return:
         """
         try:
-        c = conn.cursor()
+            c = self.conn.cursor()
             c.execute(create_table_sql)
         except Error as e:
             print(e)
 
-def create_chromosome(conn, chromosome):
+    def insert_chromosome(self, chromosome):
         """
         Create a new task
         :param conn:
@@ -40,46 +48,32 @@ def create_chromosome(conn, chromosome):
 
         sql = ''' INSERT INTO data(generation,fitness,chromosome)
                   VALUES(?,?,?) '''
-    cur = conn.cursor()
+        cur = self.conn.cursor()
         cur.execute(sql, chromosome)
-    conn.commit()
+        self.conn.commit()
         return cur.lastrowid
 
-
+    def create_data_table(self, ga):
-def main():
-    database = r"database.db"
-
-    sql_create_data_table = """ CREATE TABLE IF NOT EXISTS data (
-                                        id integer PRIMARY KEY,
-                                        generation integer NOT NULL,
-                                        fitness DOUBLE,
-                                        chromosome text
-                                    ); """
-
 
         # create a database connection
-    conn = create_connection(database)
+        self.conn = self.create_connection(ga.database_name)
 
         # create tables
-    if conn is not None:
+        if self.conn is not None:
             # create projects table
-        create_table(conn, sql_create_data_table)
+            self.create_table(ga.sql_create_data_structure)
 
             # create tasks table
             # create_table(conn, sql_create_tasks_table)
         else:
             print("Error! cannot create the database connection.")
 
-    with conn:
+        with self.conn:
 
             # Create a fake chromosome
             # Generation / Fitness / chromosome
             data1 = (0, 99, '[gene,gene,gene,gene]')
             data2 = (1, 200, '[gene,gene,gene,gene]')
             # Add chromosome to the data table inside the database
-        create_chromosome(conn, data1)
+            self.insert_chromosome(data1)
-        create_chromosome(conn, data2)
+            self.insert_chromosome(data2)
-
-
-if __name__ == '__main__':
-    main()