Repository: aeye-lab/pymovements
Pull request #671 (2024-01-31 13:48:56)

feat: add EventDataFrame.add_trial_column()
## Description
A helpful method for working on  #655
In hindsight I'm not so sure anymre if it is useful at all to have it as an EventDataFrame method

## Implemented changes

- [x] implement add_trial_column() method

## Type of change

- [x] New feature (non-breaking change which adds functionality)

## How Has This Been Tested?

- [x] single column
- [x] two columns
- [x] test for exception raising

## Checklist:

- [x] My code follows the style guidelines of this project
- [x] I have performed a self-review of my own code
- [x] I have commented my code, particularly in hard-to-understand areas
- [x] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective or that my feature works
- [x] New and existing unit tests pass locally with my changes
- [x] Any dependent changes have been merged and published in downstream modules
- [x] I have checked my code and corrected any misspellings


codecov[bot]: ## [Codecov](https://app.codecov.io/gh/aeye-lab/pymovements/pull/671?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=aeye-lab) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Comparison is base [(`8f17a97`)](https://app.codecov.io/gh/aeye-lab/pymovements/commit/8f17a970055c803f3e7b2792e2c98ffcb306cbd1?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=aeye-lab) 100.00% compared to head [(`e881d73`)](https://app.codecov.io/gh/aeye-lab/pymovements/pull/671?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=aeye-lab) 100.00%.


<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff            @@
##              main      #671   +/-   ##
=========================================
  Coverage   100.00%   100.00%           
=========================================
  Files           54        54           
  Lines         2366      2366           
  Branches       590       590           
=========================================
  Hits          2366      2366           
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/aeye-lab/pymovements/pull/671?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=aeye-lab).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=aeye-lab).


Patch:
diff --git a/src/pymovements/events/frame.py b/src/pymovements/events/frame.py
index d6c34ae..cec000d 100644
--- a/src/pymovements/events/frame.py
+++ b/src/pymovements/events/frame.py
@@ -181,6 +181,45 @@ class EventDataFrame:
         """
         self.frame = self.frame.join(event_properties, on=join_on, how='left')
 
+    def add_trial_column(
+            self,
+            column: str | list[str],
+            data: int | float | str | list[int | float | str] | None,
+    ) -> None:
+        """Add new trial columns with constant values.
+
+        Parameters
+        ----------
+        column: str | list[str]
+            The name(s) of the new trial column(s).
+        data: int | float | str | list[int | float | str] | None
+            The values to be used for filling the trial column(s). In case multiple columns are
+            provided, data must be a list of values matching the provided column order.
+        """
+        # Create trial column dictionary to iterate over in select().
+        if isinstance(column, str):
+            trial_columns = {column: data}
+        # In case a list of a single column is passed as an explicit value.
+        elif len(column) == 1 and (isinstance(data, (int, float, str) or data is None)):
+            trial_columns = {column[0]: data}
+        else:
+            if not isinstance(data, Sequence):
+                raise TypeError(
+                    'data must be passed as a list of values in case of providing multiple columns',
+                )
+            checks.check_is_length_matching(column=column, data=data)
+
+            trial_columns = dict(zip(column, data))
+
+        self.frame = self.frame.select(
+            [
+                pl.lit(column_data).alias(column_name) if not isinstance(column_data, int)
+                # Enforce Int64 columns for integers.
+                else pl.lit(column_data).alias(column_name).cast(pl.Int64)
+                for column_name, column_data in trial_columns.items()
+            ] + [pl.all()],
+        )
+
     @property
     def event_property_columns(self) -> list[str]:
         """Event property columns for this dataframe.


Test patch:
diff --git a/tests/unit/events/frame_test.py b/tests/unit/events/frame_test.py
index 8b49340..c967284 100644
--- a/tests/unit/events/frame_test.py
+++ b/tests/unit/events/frame_test.py
@@ -407,3 +407,82 @@ def test_event_dataframe_copy():
     assert events is not events_copy
     assert events.frame is not events_copy.frame
     assert_frame_equal(events.frame, events_copy.frame)
+
+
+@pytest.mark.parametrize(
+    ('events', 'kwargs', 'expected_df'),
+    [
+        pytest.param(
+            pm.EventDataFrame(name='a', onsets=[0], offsets=[1]),
+            {'column': 'trial', 'data': 1},
+            pm.EventDataFrame(
+                pl.DataFrame(
+                    {'trial': [1], 'name': 'a', 'onset': [0], 'offset': [1]},
+                ),
+            ),
+            id='single_row_trial_str',
+        ),
+        pytest.param(
+            pm.EventDataFrame(name='a', onsets=[0], offsets=[1]),
+            {'column': ['trial'], 'data': 1},
+            pm.EventDataFrame(
+                pl.DataFrame(
+                    {'trial': [1], 'name': 'a', 'onset': [0], 'offset': [1]},
+                ),
+            ),
+            id='single_row_trial_list_data_int',
+        ),
+        pytest.param(
+            pm.EventDataFrame(name='a', onsets=[0], offsets=[1]),
+            {'column': ['trial'], 'data': [1]},
+            pm.EventDataFrame(
+                pl.DataFrame(
+                    {'trial': [1], 'name': 'a', 'onset': [0], 'offset': [1]},
+                ),
+            ),
+            id='single_row_trial_list_single_identifier',
+        ),
+        pytest.param(
+            pm.EventDataFrame(name='a', onsets=[0], offsets=[1]),
+            {'column': ['group', 'trial'], 'data': ['A', 1]},
+            pm.EventDataFrame(
+                pl.DataFrame(
+                    {'group': 'A', 'trial': [1], 'name': 'a', 'onset': [0], 'offset': [1]},
+                ),
+            ),
+            id='single_row_trial_list_single_identifier',
+        ),
+        pytest.param(
+            pm.EventDataFrame(name='a', onsets=[0, 8], offsets=[1, 9]),
+            {'column': ['trial'], 'data': [1]},
+            pm.EventDataFrame(
+                pl.DataFrame(
+                    {'trial': [1, 1], 'name': ['a', 'a'], 'onset': [0, 8], 'offset': [1, 9]},
+                ),
+            ),
+            id='two_rows_trial_list_single_identifier',
+        ),
+    ],
+)
+def test_event_dataframe_add_trial_column(events, kwargs, expected_df):
+    events.add_trial_column(**kwargs)
+    assert_frame_equal(events.frame, expected_df.frame)
+
+
+@pytest.mark.parametrize(
+    ('events', 'kwargs', 'exception', 'message'),
+    [
+        pytest.param(
+            pm.EventDataFrame(name='a', onsets=[0], offsets=[1]),
+            {'column': ['group', 'trial'], 'data': 1},
+            TypeError,
+            'data must be passed as a list of values in case of providing multiple columns',
+            id='multiple_columns_data_not_list',
+        ),
+    ],
+)
+def test_event_dataframe_add_trial_column_raises_exception(events, kwargs, exception, message):
+    with pytest.raises(exception) as excinfo:
+        events.add_trial_column(**kwargs)
+
+    assert message == excinfo.value.args[0]