From 029c09e9fa98ff4be32f98ca30034407033b18e9 Mon Sep 17 00:00:00 2001
From: Koen Vossen <koen@teamtv.nl>
Date: Thu, 2 Jul 2020 10:24:27 +0200
Subject: [PATCH] Add test for additional_columns

---
 CHANGES.txt                  |  4 +++-
 kloppy/helpers.py            |  8 ++++++--
 kloppy/tests/test_helpers.py | 30 ++++++++++++++++++++++++++++++
 setup.py                     |  2 +-
 4 files changed, 40 insertions(+), 4 deletions(-)

diff --git a/CHANGES.txt b/CHANGES.txt
index 133117be..0bc81c68 100644
--- 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
\ No newline at end of file
+                      Improved ball_recovery example
+v0.6.1, 2020-07-02 -- Fix in readme (@rjtavares)
+                      Add additional_columns to to_pandas (@rjtavares)
\ No newline at end of file
diff --git a/kloppy/helpers.py b/kloppy/helpers.py
index 025fac76..cda061c7 100644
--- 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
index 650c04a2..2252e6b8 100644
--- 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
index 6391f46f..fe75d914 100644
--- 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="info@koenvossen.nl",
     url="https://github.com/PySport/kloppy",