[feat] add tracking.db loading

Author: bgallois

fastanalysis/load.py

@@ -1,4 +1,5 @@
 import pandas
+import sqlite3
 import os
 
 
@@ -8,13 +9,19 @@ class Load:
     def __init__(self, path):
         """Constructor for the Load tracking object.
 
-        :param path: Path to the tracking.txt file.
+        :param path: Path to the tracking.txt/tracking.db file.
         :type path: str
+        :type path: float
         :raises Exception
         """
         self.path = os.path.abspath(path)
         try:
-            self.tracking = pandas.read_csv(path, sep='\t')
+            if ".txt" in self.path:
+                self.tracking = pandas.read_csv(path, sep='\t')
+            else:
+                cnx = sqlite3.connect(self.path)
+                self.tracking = pandas.read_sql_query("SELECT * FROM tracking", cnx)
+                cnx.close()
         except Exception as e:
             raise e
 

fastanalysis/tests/test_load.py

@@ -10,6 +10,11 @@ def test_load_file():
     reference = pandas.read_csv("tests/tracking.txt", sep='\t')
     pandas.testing.assert_frame_equal(tracking, reference)
 
+def test_load_database():
+    """Test that file is loaded as a dataframe."""
+    tracking = load.Load("tests/tracking.db").getDataframe()
+    reference = pandas.read_csv("tests/tracking.txt", sep='\t')
+    pandas.testing.assert_frame_equal(tracking, reference)
 
 def test_load_file_error():
     """Test that wrong path lead to Exception."""
@@ -20,14 +25,16 @@ def test_load_file_error():
 
 def test_object_number():
     """Test number of objects."""
+    reference = pandas.read_csv("tests/tracking.txt", sep='\t')
     objectNumber = load.Load("tests/tracking.txt").getObjectNumber()
-    assert objectNumber == 207
+    assert objectNumber == reference.id.max() + 1
 
 
 def test_object_ids():
     """Test number of objects."""
+    reference = pandas.read_csv("tests/tracking.txt", sep='\t')
     objectNumber = load.Load("tests/tracking.txt").getIds()
-    assert objectNumber == list(range(207))
+    assert objectNumber == list(set(reference.id))
 
 
 def test_get_keys():