Merge pull request #50 from eujern/add-dataset-type

Add DatasetType.

kloppy/domain/models/common.py

@@ -1,4 +1,4 @@
-from abc import ABC
+from abc import ABC, abstractmethod
 from dataclasses import dataclass, field
 from enum import Enum, Flag
 from typing import Optional, List, Dict
@@ -224,7 +224,17 @@ class Metadata:
     provider: Provider
 
 
+class DatasetType(Enum):
+    TRACKING = "TRACKING"
+    EVENT = "EVENT"
+
+
 @dataclass
 class Dataset(ABC):
     records: List[DataRecord]
     metadata: Metadata
+
+    @property
+    @abstractmethod
+    def dataset_type(self) -> DatasetType:
+        raise NotImplementedError

kloppy/domain/models/event.py

@@ -1,11 +1,13 @@
 # Metrica Documentation https://github.com/metrica-sports/sample-data/blob/master/documentation/events-definitions.pdf
-from abc import ABC, abstractmethod, abstractproperty, ABCMeta
+from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from enum import Enum
 from typing import List, Union, Dict
 
-from .pitch import Point
+from kloppy.domain.models.common import DatasetType
+
 from .common import DataRecord, Dataset, Team, Player
+from .pitch import Point
 
 
 class ResultType(Enum):
@@ -139,6 +141,8 @@ class EventDataset(Dataset):
         Union[GenericEvent, ShotEvent, PassEvent, TakeOnEvent, CarryEvent]
     ]
 
+    dataset_type: DatasetType = DatasetType.EVENT
+
     @property
     def events(self):
         return self.records

kloppy/domain/models/tracking.py

@@ -1,7 +1,9 @@
 from dataclasses import dataclass
 from typing import List, Dict
 
-from .common import Dataset, DataRecord, Ground, Player
+from kloppy.domain.models.common import DatasetType
+
+from .common import Dataset, DataRecord, Player
 from .pitch import Point
 
 
@@ -16,6 +18,8 @@ class Frame(DataRecord):
 class TrackingDataset(Dataset):
     records: List[Frame]
 
+    dataset_type: DatasetType = DatasetType.TRACKING
+
     @property
     def frames(self):
         return self.records

kloppy/domain/services/transformers/__init__.py

@@ -109,15 +109,15 @@ def transform_event(self, event: EventType) -> EventType:
 
         return replace(event, **position_changes)
 
-    DatasetType = TypeVar("DatasetType")
+    DatasetT = TypeVar("DatasetT")
 
     @classmethod
     def transform_dataset(
         cls,
-        dataset: DatasetType,
+        dataset: DatasetT,
         to_pitch_dimensions: PitchDimensions = None,
         to_orientation: Orientation = None,
-    ) -> DatasetType:
+    ) -> DatasetT:
         if not to_pitch_dimensions and not to_orientation:
             return dataset
         elif not to_orientation:

kloppy/helpers.py

@@ -115,12 +115,12 @@ def load_metrica_json_event_data(
         )
 
 
-DatasetType = TypeVar("DatasetType")
+DatasetT = TypeVar("DatasetT")
 
 
 def transform(
-    dataset: DatasetType, to_orientation=None, to_pitch_dimensions=None
-) -> DatasetType:
+    dataset: DatasetT, to_orientation=None, to_pitch_dimensions=None
+) -> DatasetT:
     if to_orientation and isinstance(to_orientation, str):
         to_orientation = Orientation[to_orientation]
     if to_pitch_dimensions and (

kloppy/tests/test_helpers.py

@@ -28,6 +28,7 @@
     Ground,
     Player,
 )
+from kloppy.domain.models.common import DatasetType
 
 
 class TestHelpers:
@@ -40,6 +41,7 @@ def test_load_metrica_tracking_data(self):
         assert len(dataset.records) == 6
         assert len(dataset.metadata.periods) == 2
         assert dataset.metadata.provider == Provider.METRICA
+        assert dataset.dataset_type == DatasetType.TRACKING
 
     def test_load_tracab_tracking_data(self):
         base_dir = os.path.dirname(__file__)
@@ -50,6 +52,7 @@ def test_load_tracab_tracking_data(self):
         assert len(dataset.records) == 5  # only alive=True
         assert len(dataset.metadata.periods) == 2
         assert dataset.metadata.provider == Provider.TRACAB
+        assert dataset.dataset_type == DatasetType.TRACKING
 
     def _get_tracking_dataset(self):
         home_team = Team(team_id="home", name="home", ground=Ground.HOME)

kloppy/tests/test_metrica.py

@@ -8,6 +8,7 @@
     Orientation,
     Point,
 )
+from kloppy.domain.models.common import DatasetType
 
 
 class TestMetricaTracking:
@@ -28,6 +29,7 @@ def test_correct_deserialization(self):
                 }
             )
         assert dataset.metadata.provider == Provider.METRICA
+        assert dataset.dataset_type == DatasetType.TRACKING
         assert len(dataset.records) == 6
         assert len(dataset.metadata.periods) == 2
         assert dataset.metadata.orientation == Orientation.FIXED_HOME_AWAY
@@ -88,6 +90,7 @@ def test_correct_deserialization(self):
             )
 
         assert dataset.metadata.provider == Provider.METRICA
+        assert dataset.dataset_type == DatasetType.EVENT
         assert len(dataset.events) == 3620
         assert len(dataset.metadata.periods) == 2
         assert dataset.metadata.orientation is None

kloppy/tests/test_opta.py

@@ -10,6 +10,7 @@
     Position,
     Ground,
 )
+from kloppy.domain.models.common import DatasetType
 
 
 class TestOpta:
@@ -26,6 +27,7 @@ def test_correct_deserialization(self):
                 inputs={"f24_data": f24_data, "f7_data": f7_data}
             )
         assert dataset.metadata.provider == Provider.OPTA
+        assert dataset.dataset_type == DatasetType.EVENT
         assert len(dataset.events) == 17
         assert len(dataset.metadata.periods) == 2
         assert dataset.events[10].ball_owning_team == dataset.metadata.teams[1]

kloppy/tests/test_statsbomb.py

@@ -9,6 +9,7 @@
     Position,
     Provider,
 )
+from kloppy.domain.models.common import DatasetType
 
 
 class TestStatsbomb:
@@ -31,6 +32,7 @@ def test_correct_deserialization(self):
             )
 
         assert dataset.metadata.provider == Provider.STATSBOMB
+        assert dataset.dataset_type == DatasetType.EVENT
         assert len(dataset.events) == 4002
         assert len(dataset.metadata.periods) == 2
         assert (

kloppy/tests/test_tracab.py

@@ -11,6 +11,7 @@
     Team,
     Ground,
 )
+from kloppy.domain.models.common import DatasetType
 
 
 class TestTracabTracking:
@@ -29,6 +30,7 @@ def test_correct_deserialization(self):
             )
 
         assert dataset.metadata.provider == Provider.TRACAB
+        assert dataset.dataset_type == DatasetType.TRACKING
         assert len(dataset.records) == 6
         assert len(dataset.metadata.periods) == 2
         assert dataset.metadata.orientation == Orientation.FIXED_HOME_AWAY