Add test for additional_columns

PySport · Jul 2, 2020 · 029c09e · 029c09e
1 parent 5859583
commit 029c09e
Show file tree

Hide file tree

Showing 4 changed files with 40 additions and 4 deletions.
diff --git a/CHANGES.txt b/CHANGES.txt
@@ -25,4 +25,6 @@ v0.5.3, 2020-06-16 -- Add code formatting and contributing guide (@dmallory42)
 v0.6.0, 2020-06-18 -- Add Opta event serializer
                       Fix for event pattern matching for nested captures
                       Fix for event pattern matching when multiple paths can match
-                      Improved ball_recovery example
+                      Improved ball_recovery example
+v0.6.1, 2020-07-02 -- Fix in readme (@rjtavares)
+                      Add additional_columns to to_pandas (@rjtavares)
diff --git a/kloppy/helpers.py b/kloppy/helpers.py
@@ -193,7 +193,9 @@ def _event_to_pandas_row_converter(event: Event) -> Dict:
 
 
 def to_pandas(
-    dataset: Dataset, _record_converter: Callable = None, additional_columns: Dict = None
+    dataset: Dataset,
+    _record_converter: Callable = None,
+    additional_columns: Dict = None,
 ) -> "DataFrame":
     try:
         import pandas as pd
@@ -223,7 +225,9 @@ def generic_record_converter(record: Union[Frame, Event]):
 
         return row
 
-    return pd.DataFrame.from_records(map(generic_record_converter, dataset.records))
+    return pd.DataFrame.from_records(
+        map(generic_record_converter, dataset.records)
+    )
 
 
 __all__ = [

diff --git a/kloppy/tests/test_helpers.py b/kloppy/tests/test_helpers.py
@@ -125,3 +125,33 @@ def test_to_pandas(self):
         )
 
         assert_frame_equal(data_frame, expected_data_frame)
+
+    def test_to_pandas_additional_columns(self):
+        tracking_data = self._get_tracking_dataset()
+
+        data_frame = to_pandas(
+            tracking_data,
+            additional_columns={
+                "match": "test",
+                "bonus_column": lambda frame: frame.frame_id + 10,
+            },
+        )
+
+        expected_data_frame = DataFrame.from_dict(
+            {
+                "period_id": [1, 1],
+                "timestamp": [0.1, 0.2],
+                "ball_state": [None, None],
+                "ball_owning_team": [None, None],
+                "ball_x": [100, 0],
+                "ball_y": [-50, 50],
+                "match": ["test", "test"],
+                "bonus_column": [11, 12],
+                "player_home_1_x": [None, 15],
+                "player_home_1_y": [None, 35],
+                "player_away_1_x": [None, 10],
+                "player_away_1_y": [None, 20],
+            }
+        )
+
+        assert_frame_equal(data_frame, expected_data_frame)
diff --git a/setup.py b/setup.py
@@ -7,7 +7,7 @@
 
 setup(
     name="kloppy",
-    version="0.6.0",
+    version="0.6.1",
     author="Koen Vossen",
     author_email="[email protected]",
     url="https://github.com/PySport/kloppy",