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",