bids2table

180def index_dataset(
181    root: str | PathT,
182    include_subjects: str | list[str] | None = None,
183    max_workers: int | None = 0,
184    chunksize: int = 32,
185    executor_cls: type[Executor] = ProcessPoolExecutor,
186    show_progress: bool = False,
187) -> pa.Table:
188    """Index a BIDS dataset.
189
190    Args:
191        root: BIDS dataset root directory.
192        include_subjects: Glob pattern or list of patterns for matching subjects to
193            include in the index.
194        max_workers: Number of indexing processes to run in parallel. Setting
195            `max_workers=0` (the default) uses the main process only. Setting
196            `max_workers=None` starts as many workers as there are available CPUs. See
197            `concurrent.futures.ProcessPoolExecutor` for details.
198        chunksize: Number of subjects per process task. Only used for
199            `ProcessPoolExecutor` when `max_workers > 0`.
200        executor_cls: Executor class to use for parallel indexing.
201        show_progress: Show progress bar.
202
203    Returns:
204        An Arrow table index of the BIDS dataset.
205    """
206    root = as_path(root)
207
208    schema = get_arrow_schema()
209
210    dataset, _ = _get_bids_dataset(root)
211    if dataset is None:
212        _logger.warning(f"Path {root} is not a valid BIDS dataset directory.")
213        return pa.Table.from_pylist([], schema=schema)
214
215    subject_dirs = _find_bids_subject_dirs(root, include_subjects)
216    subject_dirs = sorted(subject_dirs, key=lambda p: p.name)
217    if len(subject_dirs) == 0:
218        _logger.warning(f"Path {root} contains no matching subject dirs.")
219        return pa.Table.from_pylist([], schema=schema)
220
221    func = partial(_index_bids_subject_dir, schema=schema, dataset=dataset)
222
223    tables = []
224    file_count = 0
225    for sub, table in (
226        pbar := tqdm(
227            _pmap(func, subject_dirs, max_workers, chunksize, executor_cls),
228            desc=dataset,
229            total=len(subject_dirs),
230            disable=not show_progress,
231        )
232    ):
233        file_count += len(table)
234        pbar.set_postfix(dict(sub=sub, N=_hfmt(file_count)), refresh=False)
235        tables.append(table)
236
237    # NOTE: concat_tables produces a table where each column is a ChunkedArray, with one
238    # chunk per original subject table. Is it better to keep the original chunks (one
239    # per subject) or merge using `combine_chunks`?
240    table = pa.concat_tables(tables).combine_chunks()
241    return table

244def batch_index_dataset(
245    roots: list[str | PathT],
246    max_workers: int | None = 0,
247    executor_cls: type[Executor] = ProcessPoolExecutor,
248    show_progress: bool = False,
249) -> Generator[pa.Table, None, None]:
250    """Index a batch of BIDS datasets.
251
252    Args:
253        roots: List of BIDS dataset root directories.
254        max_workers: Number of indexing processes to run in parallel. Setting
255            `max_workers=0` (the default) uses the main process only. Setting
256            `max_workers=None` starts as many workers as there are available CPUs. See
257            `concurrent.futures.ProcessPoolExecutor` for details.
258        executor_cls: Executor class to use for parallel indexing.
259        show_progress: Show progress bar.
260
261    Yields:
262        An Arrow table index for each BIDS dataset.
263    """
264    file_count = 0
265    for dataset, table in (
266        pbar := tqdm(
267            _pmap(_batch_index_func, roots, max_workers, executor_cls=executor_cls),
268            total=len(roots) if isinstance(roots, Sequence) else None,
269            disable=show_progress not in {True, "dataset"},
270        )
271    ):
272        file_count += len(table)
273        pbar.set_postfix(dict(ds=dataset, N=_hfmt(file_count)), refresh=False)
274        yield table

130def find_bids_datasets(
131    root: str | PathT,
132    exclude: str | list[str] | None = None,
133    follow_symlinks: bool = True,
134    log_frequency: int = 100,
135) -> Generator[PathT, None, None]:
136    """Find all BIDS datasets under a root directory.
137
138    Args:
139        root: Root path to begin search.
140        exclude: Glob pattern or list of patterns matching sub-directory names to
141            exclude from the search.
142        follow_symlinks: Search into symlinks that point to directories.
143
144    Yields:
145        Root paths of all BIDS datasets under `root`.
146    """
147    root = as_path(root)
148
149    dir_count = 0
150    ds_count = 0
151
152    # NOTE: Path.walk was introduced in 3.12. Otherwise, could use an older python.
153    for dirpath, dirnames, _ in root.walk(follow_symlinks=follow_symlinks):
154        dir_count += 1
155
156        if _is_bids_dataset(dirpath):
157            ds_count += 1
158            yield dirpath
159
160            # Only descend into specific sub-directories that are allowed to contain
161            # sub-datasets.
162            _filter_dirnames(dirnames, _BIDS_NESTED_PARENT_DIRNAMES)
163
164        # Filter sub-directories to descend into.
165        if exclude:
166            matches = _filter_exclude(dirnames, exclude)
167            _filter_dirnames(dirnames, matches)
168
169        if log_frequency and dir_count % log_frequency == 0:
170            _logger.info(
171                "Searched %d directories; found %d BIDS datasets.", dir_count, ds_count
172            )
173
174    if log_frequency:
175        _logger.info(
176            "Searched %d directories; found %d BIDS datasets.", dir_count, ds_count
177        )

 90def get_arrow_schema() -> pa.Schema:
 91    """Get Arrow schema of the BIDS dataset index."""
 92    entity_schema = get_bids_entity_arrow_schema()
 93    index_fields = {
 94        name: pa.field(name, cfg["dtype"], metadata=cfg["metadata"])
 95        for name, cfg in _INDEX_ARROW_FIELDS.items()
 96    }
 97    fields = [
 98        index_fields["dataset"],
 99        *entity_schema,
100        index_fields["extra_entities"],
101        index_fields["root"],
102        index_fields["path"],
103    ]
104
105    metadata = {
106        **entity_schema.metadata,
107        "bids2table_version": importlib.metadata.version(__package__),
108    }
109    schema = pa.schema(fields, metadata=metadata)
110    return schema

113def get_column_names() -> enum.StrEnum:
114    """Get an enum of the BIDS index columns."""
115    # TODO: It might be nice if the column names were statically available. One option
116    # would be to generate a static _schema.py module at install time (similar to how
117    # _version.py is generated) which defines the static default schema and column
118    # names.
119    schema = get_arrow_schema()
120    items = []
121    for f in schema:
122        name = f.metadata["name".encode()].decode()
123        items.append((name, name))
124
125    BIDSColumn = enum.StrEnum("BIDSColumn", items)
126    BIDSColumn.__doc__ = "Enum of BIDS index column names."
127    return BIDSColumn

138def parse_bids_entities(path: str | Path) -> dict[str, str]:
139    """Parse entities from BIDS file path.
140
141    Parses all BIDS filename `"{key}-{value}"` entities as well as special entities:
142    datatype, suffix, ext (extension). Does not validate entities or cast to types.
143
144    Args:
145        path: BIDS path to parse.
146
147    Returns:
148        A dict mapping BIDS entity keys to values.
149    """
150    if isinstance(path, str):
151        path = Path(path)
152    entities = {}
153
154    filename = path.name
155    parts = filename.split("_")
156
157    datatype = _parse_bids_datatype(path)
158
159    # Get suffix and extension.
160    suffix_ext = parts.pop()
161    suffix, dot, ext = suffix_ext.partition(".")
162    ext = dot + ext if ext else None
163
164    # Suffix is actually an entity, put back in list.
165    if "-" in suffix:
166        parts.append(suffix)
167        suffix = None
168
169    # Split entities, skipping any that don't contain a '-'.
170    for part in parts:
171        if "-" in part:
172            key, val = part.split("-", maxsplit=1)
173            entities[key] = val
174
175    for k, v in zip(["datatype", "suffix", "ext"], [datatype, suffix, ext]):
176        if v is not None:
177            entities[k] = v
178    return entities

197def validate_bids_entities(
198    entities: dict[str, Any],
199) -> tuple[dict[str, BIDSValue], dict[str, Any]]:
200    """Validate BIDS entities.
201
202    Validates the type and allowed values of each entity against the BIDS schema.
203
204    Args:
205        entities: dict mapping BIDS keys to unvalidated entities
206
207    Returns:
208        A tuple of `(valid_entities, extra_entities)`, where `valid_entities` is a
209            mapping of valid BIDS keys to type-casted values, and `extra_entities` a
210            mapping of any leftover entity mappings that didn't match a known entity or
211            failed validation.
212    """
213    valid_entities = {}
214    extra_entities = {}
215
216    for name, value in entities.items():
217        if name in _BIDS_NAME_ENTITY_MAP:
218            entity = _BIDS_NAME_ENTITY_MAP[name]
219            cfg = _BIDS_ENTITY_SCHEMA[entity]
220            typ = _BIDS_FORMAT_PY_TYPE_MAP[cfg["format"]]
221
222            # Cast to target type.
223            try:
224                value = typ(value)
225            except ValueError:
226                _logger.warning(
227                    f"Unable to coerce {repr(value)} to type {typ} for entity '{name}'.",
228                )
229                extra_entities[name] = value
230                continue
231
232            # Check allowed values.
233            if "enum" in cfg and value not in cfg["enum"]:
234                _logger.warning(
235                    f"Value {value} for entity '{name}' isn't one of the "
236                    f"allowed values: {cfg['enum']}.",
237                )
238                extra_entities[name] = value
239                continue
240
241            valid_entities[name] = value
242        else:
243            extra_entities[name] = value
244
245    return valid_entities, extra_entities

 78def set_bids_schema(path: str | Path | None = None) -> None:
 79    """Set the BIDS schema."""
 80    global _BIDS_SCHEMA, _BIDS_ENTITY_SCHEMA, _BIDS_NAME_ENTITY_MAP
 81    global _BIDS_ENTITY_ARROW_SCHEMA
 82
 83    schema = bidsschematools.schema.load_schema(path)
 84    entity_schema = {
 85        entity: schema.objects.entities[entity].to_dict()
 86        for entity in schema.rules.entities
 87    }
 88    # Also include special extra entities (datatype, suffix, extension).
 89    entity_schema.update(_BIDS_SPECIAL_ENTITY_SCHEMA)
 90    name_entity_map = {cfg["name"]: entity for entity, cfg in entity_schema.items()}
 91
 92    _BIDS_SCHEMA = schema
 93    _BIDS_ENTITY_SCHEMA = entity_schema
 94    _BIDS_NAME_ENTITY_MAP = name_entity_map
 95
 96    _BIDS_ENTITY_ARROW_SCHEMA = _bids_entity_arrow_schema(
 97        entity_schema,
 98        bids_version=schema["bids_version"],
 99        schema_version=schema["schema_version"],
100    )

128def get_bids_schema() -> Namespace:
129    """Get the current BIDS schema."""
130    return _BIDS_SCHEMA

133def get_bids_entity_arrow_schema() -> pa.Schema:
134    """Get the current BIDS entity schema in Arrow format."""
135    return _BIDS_ENTITY_ARROW_SCHEMA

248def format_bids_path(entities: dict[str, Any], int_format: str = "%d") -> Path:
249    """Construct a formatted BIDS path from entities dict.
250
251    Args:
252        entities: dict mapping BIDS keys to values.
253        int_format: format string for integer (index) BIDS values.
254
255    Returns:
256        A formatted `Path` instance.
257    """
258    special = {"datatype", "suffix", "ext"}
259
260    # Formatted key-value entities.
261    entities_fmt = []
262    for name, value in entities.items():
263        if name not in special:
264            if isinstance(value, int):
265                value = int_format % value
266            entities_fmt.append(f"{name}-{value}")
267    name = "_".join(entities_fmt)
268
269    # Append suffix and extension.
270    if suffix := entities.get("suffix"):
271        name += f"_{suffix}"
272    if ext := entities.get("ext"):
273        name += ext
274
275    # Prepend parent directories.
276    path = Path(name)
277    if datatype := entities.get("datatype"):
278        path = datatype / path
279    if ses := entities.get("ses"):
280        path = f"ses-{ses}" / path
281    path = f"sub-{entities['sub']}" / path
282    return path

29def cloudpathlib_is_available() -> bool:
30    """Check if cloudpathlib is available."""
31    return _CLOUDPATHLIB_AVAILABLE