diff python/storage.py @ 341:2f39c4ed0b62

first version of indicator saving to sqlite
author Nicolas Saunier <nicolas.saunier@polymtl.ca>
date Thu, 20 Jun 2013 14:31:24 -0400
parents 1046b7346886
children 4d69486869a5
line wrap: on
line diff
--- a/python/storage.py	Wed Jun 19 23:35:24 2013 -0400
+++ b/python/storage.py	Thu Jun 20 14:31:24 2013 -0400
@@ -2,8 +2,7 @@
 # -*- coding: utf-8 -*-
 '''Various utilities to save and load data'''
 
-import utils
-import moving
+import utils, moving, events
 
 __metaclass__ = type
 
@@ -16,7 +15,7 @@
 # Sqlite
 #########################
 
-def writeTrajectoriesToSqlite(objects, outFilename, trajectoryType, objectNumbers = -1):
+def saveTrajectoriesToSqlite(objects, outFilename, trajectoryType, objectNumbers = -1):
     """
     This function writers trajectories to a specified sqlite file
     @param[in] objects -> a list of trajectories
@@ -29,7 +28,7 @@
     connection = sqlite3.connect(outFilename)
     cursor = connection.cursor()
 
-    schema = "CREATE TABLE \"positions\"(trajectory_id INTEGER,frame_number INTEGER, x_coordinate REAL, y_coordinate REAL, PRIMARY KEY(trajectory_id, frame_number))"
+    schema = "CREATE TABLE IF NOT EXISTS \"positions\"(trajectory_id INTEGER,frame_number INTEGER, x_coordinate REAL, y_coordinate REAL, PRIMARY KEY(trajectory_id, frame_number))"
     cursor.execute(schema)
 
     trajectory_id = 0
@@ -205,28 +204,50 @@
     'Deletes all indicator data in db'
     pass
 
+def createInteractionTable(cursor):
+    cursor.execute('CREATE TABLE IF NOT EXISTS interactions (id INTEGER PRIMARY KEY, object_id1 INTEGER, object_id2 INTEGER, FOREIGN KEY(object_id1) REFERENCES objects(id), FOREIGN KEY(object_id2) REFERENCES objects(id))')
+
+def createIndicatorTables(cursor):
+    # cursor.execute('CREATE TABLE IF NOT EXISTS indicators (id INTEGER PRIMARY KEY, interaction_id INTEGER, indicator_type INTEGER, FOREIGN KEY(interaction_id) REFERENCES interactions(id))')
+    # cursor.execute('CREATE TABLE IF NOT EXISTS indicator_values (indicator_id INTEGER, frame_number INTEGER, value REAL, FOREIGN KEY(indicator_id) REFERENCES indicators(id), PRIMARY KEY(indicator_id, frame_number))')
+    cursor.execute('CREATE TABLE IF NOT EXISTS indicators (interaction_id INTEGER, indicator_type INTEGER, frame_number INTEGER, value REAL, FOREIGN KEY(interaction_id) REFERENCES interactions(id), PRIMARY KEY(interaction_id, indicator_type, frame_number))')
+
+def saveInteraction(cursor, interaction):
+    roadUserNumbers = list(interaction.getRoadUserNumbers())
+    cursor.execute('INSERT INTO interactions VALUES({}, {}, {})'.format(interaction.getNum(), roadUserNumbers[0], roadUserNumbers[1]))
+
 def saveInteractions(filename, interactions):
     'Saves the interactions in the table'
     import sqlite3
     connection = sqlite3.connect(filename)
     cursor = connection.cursor()
-    cursor.execute('CREATE TABLE interactions IF NOT EXISTS (id INTEGER PRIMARY KEY, object_id1 INTEGER, object_id2 INTEGER, FOREIGN KEY(object_id1) REFERENCES objects(id), FOREIGN KEY(object_id2) REFERENCES objects(id))')
-    # get the highest interaction id
-    for i in interactions:
-        cursor.execute('INSERT INTO interactions VALUES({})'.format(i.getNum())) # todo getRoadUserNumbers()
-    # CREATE TABLE IF NOT EXISTS interactions (id INTEGER PRIMARY KEY, object_id1 INTEGER, object_id2 INTEGER, FOREIGN KEY(object_id1) REFERENCES objects(id), FOREIGN KEY(object_id2) REFERENCES objects(id));
-    # CREATE TABLE IF NOT EXISTS indicators (id INTEGER PRIMARY KEY, interaction_id INTEGER, indicator_type INTEGER, FOREIGN KEY(interaction_id) REFERENCES interactions(id))
-    # CREATE TABLE IF NOT EXISTS indicator_values (indicator_id INTEGER, frame_number INTEGER, value REAL, FOREIGN KEY(indicator_id) REFERENCES indicators(id), PRIMARY KEY(indicator_id, frame_number))
-
+    createInteractionTable(cursor)
+    for inter in interactions:
+        saveInteraction(cursor, inter)
+    connection.commit()
     connection.close()
 
-def saveIndicators(filename, indicators):
+def saveIndicator(cursor, interactionNum, indicator):
+    for instant in indicator.getTimeInterval():
+        if indicator[instant]:
+            cursor.execute('INSERT INTO indicators VALUES({}, {}, {}, {})'.format(interactionNum, events.Interaction.indicatorNameToIndices[indicator.getName()], instant, indicator[instant]))
+
+def saveIndicators(filename, interactions, indicatorNames):
     'Saves the indicator values in the table'
     import sqlite3
     connection = sqlite3.connect(filename)
-    
+    cursor = connection.cursor()
+    createInteractionTable(cursor)
+    createIndicatorTables(cursor)
+    for inter in interactions:
+        saveInteraction(cursor, inter)
+        for indicatorName in indicatorNames:
+            indicator = inter.getIndicator(indicatorName)
+            if indicator != None:
+                saveIndicator(cursor, inter.getNum(), indicator)
+    connection.commit()
+    connection.close()
 
-    connection.close()
 
 #########################
 # txt files