bids2table

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

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

130def find_bids_datasets(
131    root: str | PathT,
132    exclude: str | list[str] | None = None,
133    maxdepth: int | None = None,
134) -> Generator[PathT, None, None]:
135    """Find all BIDS datasets under a root directory.
136
137    Args:
138        root: Root path to begin search.
139        exclude: Glob pattern or list of patterns matching sub-directory names to
140            exclude from the search.
141        maxdepth: Maximum depth to search.
142
143    Yields:
144        Root paths of all BIDS datasets under `root`.
145    """
146    root = as_path(root)
147
148    if isinstance(exclude, str):
149        exclude = [exclude]
150    elif exclude is None:
151        exclude = []
152    exclude = [re.compile(fnmatch.translate(pat)) for pat in exclude]
153
154    entry_count = 1
155    ds_count = 0
156
157    if _is_bids_dataset(root):
158        ds_count += 1
159        yield root
160
161    # Tuple of path, depth
162    stack = [(root, 0)]
163
164    while stack:
165        top, depth = stack.pop()
166
167        inside_bids = _is_bids_dataset(top)
168        depth += 1
169
170        for entry in top.iterdir():
171            entry_count += 1
172
173            if any(re.fullmatch(pat, entry.name) for pat in exclude):
174                continue
175
176            if _is_bids_dataset(entry):
177                ds_count += 1
178                yield entry
179
180            # Checks if we should descend into this directory.
181            # Check not reached final depth.
182            descend = maxdepth is None or depth < maxdepth
183            # Heuristic checks whether the filename looks like a (visible) directory.
184            descend = descend and not (entry.suffix or entry.name.startswith("."))
185            # Only descend into specific subdirectories of BIDS directories.
186            descend = descend and (
187                not inside_bids or entry.name in _BIDS_NESTED_PARENT_DIRNAMES
188            )
189            # Finally, check if actually a directory (which is slow so we want to
190            # short-circuit as much as possible).
191            if descend and entry.is_dir():
192                stack.append((entry, depth))

 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