|
| 1 | +from pathlib import Path |
| 2 | + |
| 3 | +import pytest |
| 4 | + |
| 5 | +from nibabies.utils.bids import _get_age_from_tsv |
| 6 | + |
| 7 | + |
| 8 | +def create_tsv(data: dict, out_file: Path) -> None: |
| 9 | + import pandas as pd |
| 10 | + |
| 11 | + pd.DataFrame(data).to_csv(out_file, index=False, sep='\t') |
| 12 | + |
| 13 | + |
| 14 | +def create_sidecar(tsv_file: Path, units) -> None: |
| 15 | + import json |
| 16 | + |
| 17 | + out_file = tsv_file.with_suffix('.json') |
| 18 | + data = {'age': {'Units': units}} |
| 19 | + out_file.write_text(json.dumps(data)) |
| 20 | + |
| 21 | + |
| 22 | +age = {'age': [4, 4, 4]} |
| 23 | +age_weeks = {'age_weeks': [4, 8, 12]} |
| 24 | +age_months = {'age_months': [3, 6, 9]} |
| 25 | +age_years = {'age_years': [1, 1, 2]} |
| 26 | + |
| 27 | + |
| 28 | +@pytest.mark.parametrize( |
| 29 | + ('idx_col', 'idx_val', 'data', 'sidecar', 'expected'), |
| 30 | + [ |
| 31 | + ('session_id', 'x1', age, False, None), |
| 32 | + ('session_id', 'x1', age, 'months', 4), |
| 33 | + ('session_id', 'x1', age, 'weeks', 1), # Convert from 4 weeks -> 1 month |
| 34 | + ('session_id', 'x1', age, ['months', 'weeks'], None), |
| 35 | + ('session_id', 'x2', age_weeks, False, 2), |
| 36 | + ('participant_id', 'x1', age_months, False, 3), |
| 37 | + ('participant_id', 'x3', age_years, False, 24), |
| 38 | + ('session_id', 'x3', {**age_months, **age}, False, 9), |
| 39 | + (None, None, age_months, False, 3), |
| 40 | + ], |
| 41 | +) |
| 42 | +def test_get_age_from_tsv(tmp_path, idx_col, idx_val, data, sidecar, expected): |
| 43 | + tsv_file = tmp_path / 'test-age-parsing.tsv' |
| 44 | + base = {} |
| 45 | + if idx_col is not None: |
| 46 | + base[idx_col] = ['x1', 'x2', 'x3'] |
| 47 | + create_tsv({**base, **data}, tsv_file) |
| 48 | + |
| 49 | + if sidecar: |
| 50 | + create_sidecar(tsv_file, sidecar) |
| 51 | + |
| 52 | + res = _get_age_from_tsv(tsv_file, idx_col, idx_val) |
| 53 | + assert res == expected |
0 commit comments