Split aoi dataframe by chosen criteria

src/pymovements/stimulus/text.py

@@ -80,6 +80,20 @@ def __init__(
         self.end_y_column = end_y_column
         self.page_column = page_column
 
+    def split_aois_by(
+            self,
+            by: str | list[str],
+    ) -> None:
+        """Split the AOI df.
+
+        Parameters
+        ----------
+        by: str | list[str]
+            Splitting criteria
+
+        """
+        self.aois = self.aois.partition_by(by=by, as_dict=False)
+
 
 def from_file(
         aoi_path: str | Path,

tests/unit/stimulus/text_test.py

@@ -226,3 +226,34 @@ def test_text_stimulus_unsupported_format():
     expected = 'unsupported file format ".pickle".Supported formats are: '\
         '[\'.csv\', \'.tsv\', \'.txt\']'
     assert msg == expected
+
+
+@pytest.mark.parametrize(
+    ('aoi_file', 'custom_read_kwargs'),
+    [
+        pytest.param(
+            'tests/files/toy_text_1_1_aoi.csv',
+            None,
+            id='toy_text_1_1_aoi',
+        ),
+        pytest.param(
+            Path('tests/files/toy_text_1_1_aoi.csv'),
+            {'separator': ','},
+            id='toy_text_1_1_aoi',
+        ),
+    ],
+)
+def test_text_stimulus_splitting(aoi_file, custom_read_kwargs):
+    aois_df = pm.stimulus.text.from_file(
+        aoi_file,
+        aoi_column='char',
+        start_x_column='top_left_x',
+        start_y_column='top_left_y',
+        width_column='width',
+        height_column='height',
+        page_column='page',
+        custom_read_kwargs=custom_read_kwargs,
+    )
+
+    aois_df.split_aois_by(by='line_idx')
+    assert len(aois_df.aois) == 2