bids2table

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

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

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

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

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