Repository: aeye-lab/pymovements
Pull request #429 (2023-05-30 13:06:08)

feat(dataset): Allow range-objects for subset argument
## Description

This allows range objects to be passed in the `subset` argument in `dataset.load()` 

## Implemented changes

The polars function `is_in()` works perfectly well with range objects, so I just needed to add this as a valid value type.

## Type of change

- [x] New feature (non-breaking change which adds functionality)
- [x] This change is or requires a documentation update

## How Has This Been Tested?

- [x] I added a single test case which passes a range object as a value

## 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/429?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=aeye-lab) Report
Patch coverage: **`100.00`**% and no project coverage change.
> Comparison is base [(`b2c721e`)](https://app.codecov.io/gh/aeye-lab/pymovements/commit/b2c721e7246b3f3eefc2265690c64a58ba04a86f?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=aeye-lab) 100.00% compared to head [(`38aa6c7`)](https://app.codecov.io/gh/aeye-lab/pymovements/pull/429?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      #429   +/-   ##
=========================================
  Coverage   100.00%   100.00%           
=========================================
  Files           56        56           
  Lines         1991      1991           
  Branches       446       446           
=========================================
  Hits          1991      1991           
```


| [Impacted Files](https://app.codecov.io/gh/aeye-lab/pymovements/pull/429?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=aeye-lab) | Coverage Δ | |
|---|---|---|
| [src/pymovements/dataset/dataset\_files.py](https://app.codecov.io/gh/aeye-lab/pymovements/pull/429?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=aeye-lab#diff-c3JjL3B5bW92ZW1lbnRzL2RhdGFzZXQvZGF0YXNldF9maWxlcy5weQ==) | `100.00% <100.00%> (ø)` | |


</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/aeye-lab/pymovements/pull/429?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=aeye-lab).   
:loudspeaker: Do you have feedback about the report comment? [Let us know in this issue](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/docs/source/tutorials/pymovements-in-10-minutes.ipynb b/docs/source/tutorials/pymovements-in-10-minutes.ipynb
index 5143931..63e0348 100644
--- a/docs/source/tutorials/pymovements-in-10-minutes.ipynb
+++ b/docs/source/tutorials/pymovements-in-10-minutes.ipynb
@@ -203,7 +203,9 @@
     "The last two columns refer to the pixel coordinates at the timestep specified by `time`.\n",
     "\n",
     "\n",
-    "We are also able to just take a subset of the data by specifying values of the fileinfo columns:\n"
+    "We are also able to just take a subset of the data by specifying values of the fileinfo columns.\n",
+    "The key refers to the column in the `fileinfo` dataframe.\n",
+    "The values in the dictionary can be of type `bool`, `int`,  `float` or `str`, but also lists and ranges \n"
    ]
   },
   {
@@ -214,8 +216,8 @@
    "outputs": [],
    "source": [
     "subset = {\n",
-    "    'text_id': [0],\n",
-    "    'page_id': [0, 1, 2],\n",
+    "    'text_id': 0,\n",
+    "    'page_id': range(3),\n",
     "}\n",
     "dataset.load(subset=subset)\n",
     "\n",
diff --git a/src/pymovements/dataset/dataset_files.py b/src/pymovements/dataset/dataset_files.py
index 9554109..245773f 100644
--- a/src/pymovements/dataset/dataset_files.py
+++ b/src/pymovements/dataset/dataset_files.py
@@ -496,12 +496,12 @@ def take_subset(
 
         if isinstance(subset_value, (bool, float, int, str)):
             column_values = [subset_value]
-        elif isinstance(subset_value, (list, tuple)):
+        elif isinstance(subset_value, (list, tuple, range)):
             column_values = subset_value
         else:
             raise TypeError(
-                f'subset value must be of type bool, float, int, str or a list of these but'
-                f' key-value pair {subset_key}: {subset_value} is of type {type(subset_value)}',
+                f'subset values must be of type bool, float, int, str, range, or list, '
+                f'but value of pair {subset_key}: {subset_value} is of type {type(subset_value)}',
             )
 
         fileinfo = fileinfo.filter(pl.col(subset_key).is_in(column_values))


Test patch:
diff --git a/tests/dataset/dataset_test.py b/tests/dataset/dataset_test.py
index 3656585..ba067ab 100644
--- a/tests/dataset/dataset_test.py
+++ b/tests/dataset/dataset_test.py
@@ -316,12 +316,17 @@ def test_load_correct_event_dfs(dataset_configuration):
         pytest.param(
             {'subject_id': 1},
             [0],
-            id='subset_key_not_in_fileinfo',
+            id='subset_int',
         ),
         pytest.param(
             {'subject_id': [1, 11, 12]},
             [0, 2, 3],
-            id='subset_key_not_in_fileinfo',
+            id='subset_list',
+        ),
+        pytest.param(
+            {'subject_id': range(3)},
+            [0, 11],
+            id='subset_range',
         ),
     ],
 )