feat: add get_possession_sequence and get_prior_possession

statsbombpy/helpers.py

@@ -170,3 +170,70 @@ def _coord(value, i):
         if z.notna().any():
             df[f"{col}_z"] = z
     return df
+
+
+def get_possession_sequence(
+    events: pd.DataFrame, event_id: str, exclude_trigger: bool = True
+) -> pd.DataFrame:
+    """
+    Return all events in the same possession sequence as the given event.
+
+    Parameters
+    ----------
+    events : pd.DataFrame
+        Events for a single match, as returned by ``sb.events()``.
+    event_id : str
+        The ``id`` of any event in the target possession sequence.
+        Typically the goal or significant event that ended the sequence.
+    exclude_trigger : bool
+        When True (default), the event identified by ``event_id`` is
+        excluded from the result. Useful for build-up analysis where the
+        trigger (e.g. a goal) is the endpoint rather than part of the
+        sequence itself.
+
+    Returns
+    -------
+    pd.DataFrame
+        Events sorted by (period, timestamp, index). Empty DataFrame if
+        ``event_id`` is not found in ``events``.
+    """
+    trigger = events[events["id"] == event_id]
+    if trigger.empty:
+        return pd.DataFrame()
+    possession_id = trigger.iloc[0]["possession"]
+    sequence = events[events["possession"] == possession_id]
+    if exclude_trigger:
+        sequence = sequence[sequence["id"] != event_id]
+    return sequence.sort_values(["period", "timestamp", "index"]).reset_index(drop=True)
+
+
+def get_prior_possession(events: pd.DataFrame, event_id: str) -> pd.DataFrame:
+    """
+    Return all events in the possession sequence immediately before the
+    given event's possession.
+
+    Useful for set piece context: when a goal is scored from a penalty,
+    free kick, or corner, the foul or challenge that earned the set piece
+    belongs to the prior possession (``possession - 1``).
+
+    Parameters
+    ----------
+    events : pd.DataFrame
+        Events for a single match, as returned by ``sb.events()``.
+    event_id : str
+        The ``id`` of any event in the target possession sequence.
+
+    Returns
+    -------
+    pd.DataFrame
+        Events sorted by (period, timestamp, index). Empty DataFrame if
+        ``event_id`` is not found or there is no prior possession.
+    """
+    trigger = events[events["id"] == event_id]
+    if trigger.empty:
+        return pd.DataFrame()
+    possession_id = trigger.iloc[0]["possession"]
+    if possession_id <= 1:
+        return pd.DataFrame()
+    prior = events[events["possession"] == possession_id - 1]
+    return prior.sort_values(["period", "timestamp", "index"]).reset_index(drop=True)

tests/statsbombpy_test/test_helpers.py

@@ -1,4 +1,10 @@
-from statsbombpy.helpers import is_relevant
+import pandas as pd
+
+from statsbombpy.helpers import (
+    get_possession_sequence,
+    get_prior_possession,
+    is_relevant,
+)
 
 EV = {
     "type": {"name": "Pass"},
@@ -33,3 +39,60 @@ def test_multiple_filters_are_anded():
 
 def test_no_filters_matches_everything():
     assert is_relevant(EV, {})
+
+
+# Synthetic match events, deliberately shuffled. "e3" and "e4" share a
+# timestamp so only the ``index`` column can order them.
+EVENTS = pd.DataFrame(
+    [
+        {"id": "e5", "index": 5, "period": 1, "timestamp": "00:00:09.000", "possession": 2},
+        {"id": "e4", "index": 4, "period": 1, "timestamp": "00:00:07.000", "possession": 2},
+        {"id": "e1", "index": 1, "period": 1, "timestamp": "00:00:00.000", "possession": 1},
+        {"id": "e3", "index": 3, "period": 1, "timestamp": "00:00:07.000", "possession": 2},
+        {"id": "e2", "index": 2, "period": 1, "timestamp": "00:00:03.000", "possession": 1},
+        {"id": "e7", "index": 7, "period": 2, "timestamp": "00:00:00.000", "possession": 3},
+        {"id": "e6", "index": 6, "period": 1, "timestamp": "00:45:01.000", "possession": 2},
+    ]
+)
+
+
+def test_possession_sequence_excludes_trigger_by_default():
+    seq = get_possession_sequence(EVENTS, "e5")
+    assert list(seq["id"]) == ["e3", "e4", "e6"]
+
+
+def test_possession_sequence_includes_trigger_when_requested():
+    seq = get_possession_sequence(EVENTS, "e5", exclude_trigger=False)
+    assert list(seq["id"]) == ["e3", "e4", "e5", "e6"]
+
+
+def test_possession_sequence_breaks_timestamp_ties_by_index():
+    seq = get_possession_sequence(EVENTS, "e6", exclude_trigger=False)
+    assert list(seq["id"][:2]) == ["e3", "e4"]
+
+
+def test_possession_sequence_sorts_by_period_first():
+    seq = get_possession_sequence(EVENTS.assign(possession=1), "e7", exclude_trigger=False)
+    assert list(seq["id"]) == ["e1", "e2", "e3", "e4", "e5", "e6", "e7"]
+
+
+def test_possession_sequence_unknown_event_id():
+    seq = get_possession_sequence(EVENTS, "not-a-real-id")
+    assert isinstance(seq, pd.DataFrame)
+    assert seq.empty
+
+
+def test_prior_possession():
+    prior = get_prior_possession(EVENTS, "e7")
+    assert list(prior["id"]) == ["e3", "e4", "e5", "e6"]
+    assert (prior["possession"] == 2).all()
+
+
+def test_prior_possession_of_first_possession_is_empty():
+    assert get_prior_possession(EVENTS, "e1").empty
+
+
+def test_prior_possession_unknown_event_id():
+    prior = get_prior_possession(EVENTS, "not-a-real-id")
+    assert isinstance(prior, pd.DataFrame)
+    assert prior.empty