[feat] add selection by keys

Author: bgallois

README.md

@@ -9,7 +9,7 @@
  ```
  git clone https://github.com/FastTrackOrg/FastAnalysis.git
  cd FastAnalysis
- pip install fastanalysis
+ pip install .
  ```
 
  ## Development

fastanalysis/load.py

@@ -29,6 +29,27 @@ def getDataframe(self):
         else:
             return self.tracking
 
+    def getKeys(self):
+        """Get the list of keys of the tracking data.
+
+        :return: Keys
+        :rtype: List
+        """
+        return self.tracking.keys().tolist()
+
+    def getDataKeys(self, keys):
+        """Get the tracking data in a pandas DataFrame format.
+
+        :raises Exception: The selected file is empty
+        :return: Tracking data
+        :rtype: DataFrame
+        """
+        if isinstance(keys, list):
+            data = self.tracking[keys]
+        else:
+            data = self.tracking[[keys]]
+        return data
+
     def getObjectNumber(self):
         """Get the total number of objects in the tracking file.
 

fastanalysis/tests/test_load.py

@@ -19,7 +19,18 @@ def test_object_number():
     """Test number of objects."""
     objectNumber = load.Load("tests/tracking.txt").getObjectNumber()
     assert objectNumber == 207
-    
+
+def test_get_keys():
+    """Test get list of keys."""
+    keys = load.Load("tests/tracking.txt").getKeys()
+    assert keys == ["xHead", "yHead", "tHead", "xTail", "yTail", "tTail", "xBody", "yBody", "tBody", "curvature", "areaBody", "perimeterBody", "headMajorAxisLength", "headMinorAxisLength", "headExcentricity", "tailMajorAxisLength", "tailMinorAxisLength", "tailExcentricity", "bodyMajorAxisLength", "bodyMinorAxisLength", "bodyExcentricity", "imageNumber", "id"]
+
+def test_data_keys():
+    """Test get data from list of keys."""
+    reference = pandas.read_csv("tests/tracking.txt", sep='\t')
+    pandas.testing.assert_frame_equal(load.Load("tests/tracking.txt").getDataKeys(["yHead", "tHead"]), reference[["yHead", "tHead"]])
+    pandas.testing.assert_frame_equal(load.Load("tests/tracking.txt").getDataKeys("yHead"), reference[["yHead"]])
+
 def test_get_objects():
     """Test get the data for an object"""
     reference = pandas.read_csv("tests/tracking.txt", sep='\t')