bids2table

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

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

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

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

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

11def load_bids_metadata(path: str | PathT, inherit: bool = True) -> dict[str, Any]:
12    """Load the full JSON sidecar metadata for a BIDS file.
13
14    Sidecar files are loaded according to the inheritance principle in top-down order.
15
16    Args:
17        path: BIDS file path
18        inherit: Load the full metadata according to inheritance. Otherwise, load only
19            the first JSON sidecar found in the bottom-up search.
20
21    Returns:
22        A sidecar metadata dictionary.
23    """
24    path = as_path(path)
25    entities = _cache_parse_bids_entities(path)
26    query = dict(entities, ext=".json")
27
28    metadata = {}
29
30    parent = path.parent
31    if inherit:
32        sidecars = reversed(list(_find_bids_parents(parent, query)))
33    else:
34        sidecars = [next(_find_bids_parents(parent, query))]
35
36    for path in sidecars:
37        try:
38            data = _load_json(path)
39            metadata.update(data)
40        except (json.JSONDecodeError, TypeError):
41            continue
42    return metadata

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