voxkit.storage.alignments

  1"""Specialized CRUD operations for managing dataset alignments within the VoxKit storage system.
  2
  3Directory Structure
  4-------------------
  5Each dataset alignment follows a hierarchical structure:
  6
  7    dataset_id/
  8    ├── alignments/
  9    │   ├── alignment_id_1/
 10    │   │   ├── textgrids/             # Directory for TextGrid files
 11    │   │   └── voxkit_alignment.json  # Alignment metadata
 12    │   ├── alignment_id_2/
 13    │   │   └── ...
 14    │   └── ...
 15    └── ...
 16
 17API
 18---
 19- **create_alignment**: Create a new alignment entry in storage
 20- **create_hand_alignment**: Create a new hand-annotated alignment entry in storage
 21- **create_corrected_alignment**: Create a new, fully-owned alignment for correcting
 22  another alignment's boundaries, without ever mutating the source
 23- **get_alignment_metadata**: Retrieve metadata for a specific alignment
 24- **get_alignment_type**: Return an alignment's provenance (automatic/hand/corrected)
 25- **update_alignment**: Update the status or details of an existing alignment
 26- **list_alignments**: List all alignments for a given dataset
 27- **delete_alignment**: Remove an alignment from storage
 28
 29Notes
 30-----
 31- All paths are managed using pathlib
 32- Engine-specific branching may be necessary to bridge different alignment formats
 33- Error handling only exposes user-friendly messages
 34- Alignment IDs are generated using unique timestamps with microsecond precision
 35- TextGrid paths depend on whether the dataset is cached locally or not
 36"""
 37
 38import json
 39import logging
 40import os
 41import shutil
 42from pathlib import Path
 43from typing import List, Literal, NotRequired, Tuple, TypeAlias, TypedDict
 44
 45from .constants import ALIGNMENTS_ROOT, SUPERSET_AUDIO_EXTENSIONS
 46from .datasets import _get_dataset_root, get_dataset_metadata
 47from .models import ModelMetadata, get_model_metadata
 48from .utils import generate_unique_id, readable_from_unique_id
 49
 50logger = logging.getLogger(__name__)
 51
 52HAND_ALIGNMENT_SENTINEL = "hand"
 53"""Sentinel value used for engine_id/model id on manually-created (hand) alignments."""
 54
 55CORRECTED_ALIGNMENT_SENTINEL = "corrected"
 56"""Sentinel value used for engine_id/model id on boundary-corrected alignments."""
 57
 58AlignmentStatus: TypeAlias = Literal["pending", "completed", "failed"]
 59"""Status of an alignment operation.
 60
 61Values:
 62    pending: Alignment has been created but not yet processed.
 63    completed: Alignment has been successfully completed.
 64    failed: Alignment processing failed.
 65"""
 66
 67AlignmentType: TypeAlias = Literal["automatic", "hand", "corrected"]
 68"""How an alignment's TextGrids were produced -- shown as its own column in
 69GUI alignment dropdowns, kept separate from ``engine_id`` so a corrected
 70alignment can still show which real engine (mfa, w2tg, ...) it traces back
 71to, rather than losing that in favor of a "corrected" placeholder."""
 72
 73
 74class AlignmentMetadata(TypedDict):
 75    """Alignment metadata structure.
 76
 77    Attributes:
 78        id: Unique identifier (timestamp with microsecond precision).
 79        engine_id: Identifier of the alignment engine used (e.g., "mfa").
 80        model_metadata: Metadata of the model used for alignment.
 81        local: Whether TextGrid files are stored locally (cached) or at original path.
 82        alignment_date: Human-readable alignment creation timestamp.
 83        status: Current status of the alignment operation.
 84        tg_path: Path to the directory containing TextGrid output files.
 85        source_alignment_id: For corrected alignments, the id of the alignment
 86            they were corrected from. Absent on all other alignment types.
 87        alignment_type: How the alignment was produced ("automatic"/"hand"/
 88            "corrected"). Absent on alignments created before this field
 89            existed -- use ``get_alignment_type()`` rather than reading this
 90            key directly, since that also handles the fallback.
 91    """
 92
 93    id: str
 94    engine_id: str
 95    model_metadata: ModelMetadata
 96    local: bool
 97    alignment_date: str
 98    status: AlignmentStatus
 99    tg_path: str
100    source_alignment_id: NotRequired[str]
101    alignment_type: NotRequired[AlignmentType]
102
103
104def get_alignment_type(meta: AlignmentMetadata) -> AlignmentType:
105    """Return an alignment's type, inferring it for alignments that predate this field.
106
107    Older alignments (including ones created by an earlier version of
108    ``create_corrected_alignment`` that used the "corrected" sentinel as its
109    own ``engine_id``) won't have ``alignment_type`` in their stored JSON --
110    fall back to the engine_id sentinels in that case.
111    """
112    if "alignment_type" in meta:
113        return meta["alignment_type"]
114    if meta["engine_id"] == HAND_ALIGNMENT_SENTINEL:
115        return "hand"
116    if meta["engine_id"] == CORRECTED_ALIGNMENT_SENTINEL:
117        return "corrected"
118    return "automatic"
119
120
121def _get_alignments_root(dataset_id: str) -> Path | None:
122    """Get the root directory for storing alignments for a given dataset.
123
124    Args:
125        dataset_id: Identifier of the dataset
126
127    Returns:
128        Path to the alignments root directory or None if dataset not found
129    """
130    dataset_root = _get_dataset_root(dataset_id)
131    if dataset_root:
132        alignments_root = dataset_root / ALIGNMENTS_ROOT
133        alignments_root.mkdir(parents=False, exist_ok=True)
134        return alignments_root
135
136    return None
137
138
139def _get_alignment_root(dataset_id: str, alignment_id: str) -> Path | None:
140    """Get the root directory for a specific alignment by ID.
141
142    Args:
143        dataset_id: Identifier of the dataset containing the alignment
144        alignment_id: Identifier of the alignment
145
146    Returns:
147        Path to the alignment root directory or None if not found
148    """
149    alignments_root = _get_alignments_root(dataset_id)
150    if alignments_root:
151        alignment_root = alignments_root / alignment_id
152        if alignment_root.exists():
153            return alignment_root
154    return None
155
156
157def create_alignment(
158    dataset_id: str, engine_id: str, model_id: str
159) -> tuple[Literal[True], AlignmentMetadata] | tuple[Literal[False], str]:
160    """Create a new alignment entry in the storage.
161
162    Creates an alignment directory, sets up TextGrid output location based on whether
163    the dataset is cached, generates metadata, and initializes the alignment with
164    "pending" status.
165
166    Args:
167        dataset_id: Identifier of the dataset to align
168        engine_id: Identifier of the alignment engine
169        model_id: Identifier of the alignment model to use
170
171    Returns:
172        Tuple of (True, AlignmentMetadata) on success or (False, error_message) on failure
173
174    Raises:
175        FileNotFoundError: If model or dataset is not found
176        Exception: If directory creation or metadata writing fails
177
178    Notes:
179        - For cached datasets, TextGrid output is stored in the alignment directory
180        - For non-cached datasets, TextGrid output is stored in the original dataset path
181        - Automatically cleans up on failure
182    """
183    # Fetch model metadata
184    model_metadata = get_model_metadata(engine_id, model_id)
185    if not model_metadata:
186        return False, f"Model '{model_id}' for engine '{engine_id}' not found"
187
188    # Fetch dataset metadata
189    dataset_metadata = get_dataset_metadata(dataset_id)
190    if not dataset_metadata:
191        return False, f"Dataset '{dataset_id}' not found"
192
193    # Fetch alignment root
194    alignments_root = _get_alignments_root(dataset_id)
195    if not alignments_root:
196        return False, f"Dataset '{dataset_id}' not found"
197
198    # Create alignment directory
199    now = generate_unique_id()
200    alignment_date = readable_from_unique_id(now)
201    alignment_root = alignments_root / now
202
203    alignment_root.mkdir(parents=False, exist_ok=False)
204
205    try:
206        local = dataset_metadata["cached"]
207        tg_path = alignment_root / "textgrids"
208        tg_path.mkdir(parents=False, exist_ok=True)
209
210        metadata = AlignmentMetadata(
211            id=now,
212            engine_id=engine_id,
213            model_metadata=model_metadata,
214            local=local,
215            tg_path=str(tg_path),
216            alignment_date=alignment_date,
217            status="pending",
218            alignment_type="automatic",
219        )
220
221        # Fetch model metadata
222        metadata_path = alignment_root / "voxkit_alignment.json"
223
224        with open(metadata_path, "w", encoding="utf-8") as f:
225            json.dump(metadata, f, indent=4)
226
227        return True, metadata
228
229    except Exception as e:
230        # Clean up partially created directory
231        if os.path.exists(alignment_root):
232            shutil.rmtree(alignment_root, ignore_errors=True)
233        return False, f"Failed to create alignment metadata: {str(e)}"
234
235
236_AUDIO_EXTS = SUPERSET_AUDIO_EXTENSIONS
237
238
239def validate_hand_alignments(dataset_path: Path, hand_path: Path) -> Tuple[bool, str]:
240    """Validate that a hand-alignments directory matches a dataset's speaker/audio layout.
241
242    Expects the hand-alignments directory to mirror the dataset: one subdirectory
243    per speaker, containing a ``.TextGrid`` file for every audio file in the
244    corresponding dataset speaker directory (matched by stem).
245
246    Args:
247        dataset_path: Path to the dataset root (containing speaker subdirectories)
248        hand_path: Path to the hand-annotated TextGrid root
249
250    Returns:
251        Tuple of (True, "...") if valid, or (False, error_message) if not
252    """
253    if not isinstance(dataset_path, Path):
254        dataset_path = Path(dataset_path)
255    if not isinstance(hand_path, Path):
256        hand_path = Path(hand_path)
257
258    if not hand_path.exists() or not hand_path.is_dir():
259        return False, f"Hand alignments path '{hand_path}' is not an existing directory."
260
261    dataset_speakers = {
262        d.name for d in dataset_path.iterdir() if d.is_dir() and not d.name.startswith(".")
263    }
264    hand_speakers = {
265        d.name for d in hand_path.iterdir() if d.is_dir() and not d.name.startswith(".")
266    }
267
268    missing_speakers = dataset_speakers - hand_speakers
269    if missing_speakers:
270        return (
271            False,
272            f"Hand alignments missing speaker directories: {', '.join(sorted(missing_speakers))}",
273        )
274
275    for speaker in sorted(dataset_speakers):
276        audio_stems = {
277            f.stem for f in (dataset_path / speaker).iterdir() if f.suffix.lower() in _AUDIO_EXTS
278        }
279        tg_stems = {
280            f.stem for f in (hand_path / speaker).iterdir() if f.suffix.lower() == ".textgrid"
281        }
282        missing = audio_stems - tg_stems
283        if missing:
284            return (
285                False,
286                f"Speaker '{speaker}' is missing TextGrid files for: {', '.join(sorted(missing))}",
287            )
288
289    return True, "Hand alignments match dataset layout."
290
291
292def create_hand_alignment(
293    dataset_id: str,
294    tg_path: str | None = None,
295) -> tuple[Literal[True], AlignmentMetadata] | tuple[Literal[False], str]:
296    """Create a new hand-annotated alignment entry in storage.
297
298    Mirrors `create_alignment` but skips engine/model lookup — engine_id and the
299    model metadata fields are filled with the `HAND_ALIGNMENT_SENTINEL` value.
300    Starts in "completed" status since there is no processing step.
301
302    Args:
303        dataset_id: Identifier of the dataset to align
304        tg_path: Optional existing directory of hand-annotated TextGrids. When
305            provided, it is recorded as-is and the alignment is marked non-local.
306            When omitted, falls back to the same tg_path resolution used by
307            `create_alignment`.
308
309    Returns:
310        Tuple of (True, AlignmentMetadata) on success or (False, error_message) on failure
311    """
312    dataset_metadata = get_dataset_metadata(dataset_id)
313    if not dataset_metadata:
314        return False, f"Dataset '{dataset_id}' not found"
315
316    alignments_root = _get_alignments_root(dataset_id)
317    if not alignments_root:
318        return False, f"Dataset '{dataset_id}' not found"
319
320    now = generate_unique_id()
321    alignment_date = readable_from_unique_id(now)
322    alignment_root = alignments_root / now
323
324    alignment_root.mkdir(parents=False, exist_ok=False)
325
326    try:
327        if tg_path is not None:
328            local = False
329            resolved_tg_path = Path(tg_path)
330        else:
331            local = dataset_metadata["cached"]
332            if bool(local) is False:
333                resolved_tg_path = Path(dataset_metadata["original_path"]) / "textgrids"
334            else:
335                resolved_tg_path = alignment_root / "textgrids"
336            resolved_tg_path.mkdir(parents=False, exist_ok=True)
337
338        model_metadata = ModelMetadata(
339            name=HAND_ALIGNMENT_SENTINEL,
340            engine_id=HAND_ALIGNMENT_SENTINEL,
341            model_path="",  # type: ignore[typeddict-item]
342            data_path="",  # type: ignore[typeddict-item]
343            eval_path="",  # type: ignore[typeddict-item]
344            train_path="",  # type: ignore[typeddict-item]
345            download_date=alignment_date,
346            id=HAND_ALIGNMENT_SENTINEL,
347        )
348
349        metadata = AlignmentMetadata(
350            id=now,
351            engine_id=HAND_ALIGNMENT_SENTINEL,
352            model_metadata=model_metadata,
353            local=local,
354            tg_path=str(resolved_tg_path),
355            alignment_date=alignment_date,
356            status="completed",
357            alignment_type="hand",
358        )
359
360        metadata_path = alignment_root / "voxkit_alignment.json"
361        with open(metadata_path, "w", encoding="utf-8") as f:
362            json.dump(metadata, f, indent=4)
363
364        return True, metadata
365
366    except Exception as e:
367        if os.path.exists(alignment_root):
368            shutil.rmtree(alignment_root, ignore_errors=True)
369        return False, f"Failed to create hand alignment metadata: {str(e)}"
370
371
372def create_corrected_alignment(
373    dataset_id: str,
374    source_alignment_id: str,
375    engine_id: str | None = None,
376    alignment_type: str = "corrected",
377) -> tuple[Literal[True], AlignmentMetadata] | tuple[Literal[False], str]:
378    """Create a new, fully-owned alignment for hand-correcting a source alignment's boundaries.
379
380    Unlike `create_hand_alignment`, this unconditionally creates and owns its own
381    `textgrids` directory (`local=True`), regardless of the source dataset's
382    `cached` flag -- `create_hand_alignment`'s non-cached branch points `tg_path`
383    at the *original dataset directory*, which is exactly the overwrite risk this
384    function exists to avoid. The full source TextGrid set is baseline-copied in
385    immediately, so files the user never touches stay byte-identical and the
386    corrected alignment is self-contained from the start. The source alignment's
387    own TextGrids are never modified.
388
389    Args:
390        dataset_id: Identifier of the dataset
391        source_alignment_id: Identifier of the alignment being corrected
392        engine_id: Value to store/display as this alignment's Engine. Defaults
393            to the source alignment's own engine_id (i.e. which engine actually
394            produced the underlying TextGrids) if omitted, but the GUI lets a
395            user override it -- e.g. to tag who corrected it, or distinguish
396            multiple correction passes.
397        alignment_type: Value to store/display as this alignment's Type
398            (normally "automatic"/"hand"/"corrected", but this is a free-form
399            string so a user can customize it, e.g. "corrected-v2").
400
401    Returns:
402        Tuple of (True, AlignmentMetadata) on success or (False, error_message) on failure
403    """
404    dataset_metadata = get_dataset_metadata(dataset_id)
405    if not dataset_metadata:
406        return False, f"Dataset '{dataset_id}' not found"
407
408    source_metadata = get_alignment_metadata(dataset_id, source_alignment_id)
409    if not source_metadata:
410        return False, f"Source alignment '{source_alignment_id}' not found"
411
412    alignments_root = _get_alignments_root(dataset_id)
413    if not alignments_root:
414        return False, f"Dataset '{dataset_id}' not found"
415
416    now = generate_unique_id()
417    alignment_date = readable_from_unique_id(now)
418    alignment_root = alignments_root / now
419
420    alignment_root.mkdir(parents=False, exist_ok=False)
421
422    try:
423        tg_path = alignment_root / "textgrids"
424        tg_path.mkdir(parents=False, exist_ok=True)
425
426        source_tg_root = Path(source_metadata["tg_path"])
427        if source_tg_root.exists():
428            shutil.copytree(source_tg_root, tg_path, dirs_exist_ok=True)
429
430        metadata = AlignmentMetadata(
431            id=now,
432            engine_id=engine_id or source_metadata["engine_id"],
433            model_metadata=source_metadata["model_metadata"],
434            local=True,
435            tg_path=str(tg_path),
436            alignment_date=alignment_date,
437            status="completed",
438            source_alignment_id=source_alignment_id,
439            alignment_type=alignment_type,  # type: ignore[typeddict-item]
440        )
441
442        metadata_path = alignment_root / "voxkit_alignment.json"
443        with open(metadata_path, "w", encoding="utf-8") as f:
444            json.dump(metadata, f, indent=4)
445
446        return True, metadata
447
448    except Exception as e:
449        if os.path.exists(alignment_root):
450            shutil.rmtree(alignment_root, ignore_errors=True)
451        return False, f"Failed to create corrected alignment metadata: {str(e)}"
452
453
454def get_alignment_metadata(dataset_id: str, alignment_id: str) -> AlignmentMetadata | None:
455    """Get the metadata for a specific alignment by ID.
456
457    Retrieves the alignment metadata from the voxkit_alignment.json file in the
458    alignment's directory. Normalizes status values to lowercase for consistency.
459
460    Args:
461        dataset_id: Identifier of the dataset containing the alignment
462        alignment_id: Identifier of the alignment
463
464    Returns:
465        AlignmentMetadata dictionary or None if not found
466
467    Raises:
468        Exception: If metadata file cannot be loaded or parsed
469        JSONDecodeError: If the metadata file is malformed
470    """
471    alignment_root = _get_alignment_root(dataset_id, alignment_id)
472    if not alignment_root:
473        return None
474
475    metadata_path = alignment_root / "voxkit_alignment.json"
476
477    try:
478        with open(metadata_path, "r", encoding="utf-8") as f:
479            metadata: AlignmentMetadata = json.load(f)
480            # Normalize status to lowercase for consistency
481            if "status" in metadata:
482                status_lower = metadata["status"].lower()
483                # Cast to the correct literal type
484                metadata["status"] = status_lower  # type: ignore[typeddict-item]
485            return metadata
486    except Exception as e:
487        logger.exception("Failed to load alignment metadata from '%s'", metadata_path)
488        raise e
489
490
491def update_alignment(dataset_id: str, alignment_id: str, updates: dict) -> Tuple[bool, str]:
492    """Update the status or metadata of an alignment.
493
494    Updates fields in the alignment's metadata file. Only fields present in the
495    metadata are updated. Status values are automatically normalized to lowercase.
496
497    Args:
498        dataset_id: Identifier of the dataset containing the alignment
499        alignment_id: Identifier of the alignment to update
500        updates: Dictionary of updates to apply to the alignment metadata
501
502    Returns:
503        Tuple of (True, success_message) on success or (False, error_message) on failure
504
505    Raises:
506        FileNotFoundError: If the alignment is not found
507        Exception: If metadata file cannot be read or written
508
509    Notes:
510        - Status values are normalized to lowercase for consistency
511        - Commonly updated fields include "status" for tracking progress
512    """
513    alignment_root = _get_alignment_root(dataset_id, alignment_id)
514    if not alignment_root:
515        return False, f"Alignment '{alignment_id}' for dataset '{dataset_id}' not found"
516
517    metadata_path = alignment_root / "voxkit_alignment.json"
518
519    try:
520        with open(metadata_path, "r", encoding="utf-8") as f:
521            metadata = json.load(f)
522
523        # Update fields
524        for key, value in updates.items():
525            if key in metadata:
526                # Normalize status values to lowercase
527                if key == "status" and isinstance(value, str):
528                    value = value.lower()
529                metadata[key] = value
530
531        with open(metadata_path, "w", encoding="utf-8") as f:
532            json.dump(metadata, f, indent=4)
533
534        return True, "Alignment metadata updated successfully."
535
536    except Exception as e:
537        return False, f"Failed to update alignment metadata: {str(e)}"
538
539
540def list_alignments(dataset_id: str) -> List[AlignmentMetadata]:
541    """List all alignment metadata for a given dataset.
542
543    Scans the dataset's alignments directory and collects metadata from all
544    subdirectories containing valid voxkit_alignment.json files. Normalizes
545    status values to lowercase for consistency.
546
547    Args:
548        dataset_id: Identifier of the dataset to list alignments for
549
550    Returns:
551        List of AlignmentMetadata dictionaries (empty list if none found)
552
553    Notes:
554        - Skips directories with invalid or missing metadata files
555        - Returns empty list if dataset not found
556        - Status values are normalized to lowercase
557    """
558    alignments_root = _get_alignments_root(dataset_id)
559    if not alignments_root:
560        return []
561
562    alignments_found = []
563    for dir in alignments_root.iterdir():
564        if dir.is_dir():
565            metadata_path = dir / "voxkit_alignment.json"
566            if metadata_path.exists():
567                try:
568                    with open(metadata_path, "r", encoding="utf-8") as f:
569                        metadata = json.load(f)
570                        # Normalize status to lowercase for consistency
571                        if "status" in metadata:
572                            metadata["status"] = metadata["status"].lower()
573                        alignments_found.append(metadata)
574                except Exception:
575                    logger.exception("Failed to load alignment metadata from '%s'", metadata_path)
576
577    return alignments_found
578
579
580def delete_alignment(dataset_id: str, alignment_id: str) -> Tuple[bool, str]:
581    """Delete an alignment given its dataset ID and alignment ID.
582
583    Permanently removes the alignment directory and all its contents, including
584    metadata and TextGrid files (if stored locally).
585
586    Args:
587        dataset_id: Identifier of the dataset containing the alignment
588        alignment_id: Identifier of the alignment to delete
589
590    Returns:
591        Tuple of (True, success_message) on success or (False, error_message) on failure
592
593    Raises:
594        Exception: If the directory cannot be removed
595
596    Notes:
597        - This operation is irreversible
598        - Only removes TextGrid files if they are stored locally (cached datasets)
599        - TextGrid files in the original dataset path are not removed
600    """
601    alignment_root = _get_alignment_root(dataset_id, alignment_id)
602    if not alignment_root:
603        return False, f"Alignment '{alignment_id}' for dataset '{dataset_id}' not found"
604
605    try:
606        shutil.rmtree(alignment_root)
607        return True, f"Alignment '{alignment_id}' deleted successfully."
608    except Exception as e:
609        return False, f"Failed to delete alignment '{alignment_id}': {str(e)}"