voxkit.storage.datasets

  1"""Specialized CRUD operations for managing datasets within the VoxKit storage system.
  2
  3Directory Structure
  4-------------------
  5Each dataset follows a hierarchical structure:
  6
  7    my_dataset/
  8    ├── voxkit_dataset.json       # Dataset metadata
  9    ├── alignments/               # Alignment outputs storage
 10    └── cache/                    # Optional cached copy of dataset
 11        ├── speaker_001/
 12        │   ├── audio_001.wav
 13        │   ├── audio_001.lab
 14        │   └── ...
 15        └── speaker_002/
 16            └── ...
 17
 18API
 19---
 20- **create_dataset**: Create a new dataset with metadata and directories
 21- **get_dataset_metadata**: Retrieve metadata for a specific dataset
 22- **list_datasets_metadata**: List all existing datasets
 23- **update_dataset_metadata**: Update metadata fields for a specific dataset
 24- **delete_dataset**: Delete a registered dataset and its metadata
 25- **export_dataset**: Export a dataset to a specified output path
 26- **import_dataset**: Import an existing dataset into VoxKit storage
 27- **validate_dataset**: Validate dataset structure and organization
 28
 29Notes
 30-----
 31- All dataset IDs are unique timestamps with microsecond precision
 32- Failed operations automatically clean up partial changes
 33- Dataset validation occurs before creation to prevent invalid data
 34- Cached datasets are copied for faster access during operations
 35- The `transcribed` flag indicates presence of transcription files
 36- Importing datasets adjusts metadata and validates structure
 37"""
 38
 39import csv
 40import json
 41import logging
 42import os
 43import shutil
 44from pathlib import Path
 45from typing import Any, List, Literal, Tuple, TypedDict
 46
 47from voxkit.storage.constants import ALIGNMENTS_ROOT, DATASETS_ROOT, SUPERSET_AUDIO_EXTENSIONS
 48from voxkit.storage.utils import generate_unique_id, get_storage_root, readable_from_unique_id
 49
 50logger = logging.getLogger(__name__)
 51
 52
 53class DatasetMetadata(TypedDict):
 54    """Dataset metadata structure.
 55
 56    Attributes:
 57        name: Human-readable name of the dataset.
 58        id: Unique identifier (timestamp with microsecond precision).
 59        description: Description of the dataset contents and purpose.
 60        original_path: Original file system path to the dataset.
 61        cached: Whether the dataset is cached in VoxKit storage.
 62        anonymize: Whether speaker identities should be anonymized.
 63        transcribed: Whether the dataset includes transcription files.
 64        registration_date: Human-readable registration timestamp.
 65        hand_alignments_path: Optional path to pre-existing hand-annotated TextGrids.
 66    """
 67
 68    name: str
 69    id: str
 70    description: str
 71    original_path: str
 72    cached: bool
 73    anonymize: bool
 74    transcribed: bool
 75    registration_date: str
 76    hand_alignments_path: str | None
 77
 78
 79def _get_datasets_root() -> Path:
 80    """Get the root directory for datasets storage.
 81
 82    Returns:
 83        Path to datasets storage root directory
 84    """
 85    root = get_storage_root() / DATASETS_ROOT
 86    root.mkdir(parents=False, exist_ok=True)
 87    return root
 88
 89
 90def _get_dataset_root(dataset_id: str) -> Path | None:
 91    """Get the root directory for a specific dataset by ID.
 92
 93    Args:
 94        dataset_id: Identifier of the dataset
 95
 96    Returns:
 97        Path to dataset root directory or None if not found
 98    """
 99    datasets_root = _get_datasets_root()
100    if datasets_root and dataset_id:
101        dataset_root = datasets_root / dataset_id
102        if dataset_root.exists():
103            return dataset_root
104    return None
105
106
107def get_dataset_data_path(meta: DatasetMetadata) -> Path | None:
108    """Return the directory containing the dataset's speaker subdirs.
109
110    For cached datasets this is ``<dataset_root>/cache``; for non-cached
111    datasets it is the original on-disk path recorded in metadata.
112    """
113    if meta.get("cached"):
114        root = _get_dataset_root(meta["id"])
115        if root is None:
116            return None
117        return root / "cache"
118    return Path(meta["original_path"])
119
120
121def _get_dataset_metadata(dataset_root: Path) -> DatasetMetadata | None:
122    """Load dataset metadata from the given dataset root directory.
123
124    Args:
125        dataset_root: Path to the dataset root directory
126
127    Returns:
128        Dataset metadata dictionary or None if not found or invalid
129    """
130    try:
131        metadata_path = dataset_root / "voxkit_dataset.json"
132        if not metadata_path.exists():
133            return None
134        with open(metadata_path, "r", encoding="utf-8") as f:
135            result: DatasetMetadata = json.load(f)
136            return result
137    except Exception:
138        return None
139
140
141def create_dataset(
142    name: str,
143    description: str,
144    original_path: str,
145    cached: bool,
146    anonymize: bool,
147    transcribed: bool = False,
148    analysis_data: list[dict[str, Any]] | None = None,
149    analysis_method: str | None = None,
150    hand_alignments_path: str | None = None,
151) -> tuple[Literal[True], DatasetMetadata] | tuple[Literal[False], str]:
152    """Create a dataset metadata dictionary and create necessary directories.
153
154    Validates the dataset structure, creates a unique ID, sets up the directory
155    hierarchy (dataset root and alignments subdirectory), writes metadata to JSON,
156    optionally caches the dataset, and optionally saves analysis results to CSV.
157
158    Args:
159        name: Name of the dataset
160        description: Description of the dataset
161        original_path: Original path to the dataset
162        cached: Whether to copy the dataset into VoxKit storage
163        anonymize: Whether the dataset should be anonymized
164        transcribed: Whether the dataset includes transcription files
165        analysis_data: Optional list of analysis result dictionaries to save as CSV
166        analysis_method: Optional name of the analysis method (used for CSV filename)
167        hand_alignments_path: Optional path to pre-existing hand-annotated TextGrids
168
169    Returns:
170        Tuple of (True, DatasetMetadata) on success or (False, error_message) on failure
171
172    Raises:
173        FileExistsError: If a dataset with the generated ID already exists
174        Exception: If directory creation, metadata writing, or caching fails
175
176    Notes:
177        - Automatically validates dataset structure before creation
178        - Cleans up partially created directories on failure
179        - Cached datasets are copied with shutil.copytree
180        - If analysis_data is provided, saves to {analysis_method}_summary.csv
181    """
182    # Validate dataset structure
183    valid, msg = validate_dataset(Path(original_path), transcribed=transcribed)
184    if not valid:
185        return False, msg
186
187    now = generate_unique_id()
188
189    try:
190        humannow = readable_from_unique_id(now)
191        metadata = DatasetMetadata(
192            name=name,
193            id=now,
194            description=description,
195            original_path=str(original_path),
196            cached=cached,
197            anonymize=anonymize,
198            transcribed=transcribed,
199            registration_date=humannow,
200            hand_alignments_path=hand_alignments_path,
201        )
202
203        # Create dataset directory
204        dataset_dir = _get_datasets_root() / metadata["id"]
205        if dataset_dir.exists():
206            raise FileExistsError(f"Dataset with ID '{metadata['id']}' already exists.")
207        dataset_dir.mkdir(parents=False, exist_ok=False)
208
209        # Create dataset/alignments directory
210        alignments_dir = dataset_dir / ALIGNMENTS_ROOT
211        alignments_dir.mkdir(parents=False, exist_ok=False)
212        metadata_path = dataset_dir / "voxkit_dataset.json"
213        with open(metadata_path, "w", encoding="utf-8") as f:
214            json.dump(metadata, f, indent=2)
215
216        # Cache the dataset if requested
217        if cached:
218            cache_dir = dataset_dir / "cache"
219            cache_dir.mkdir(parents=False, exist_ok=False)
220            shutil.copytree(original_path, cache_dir, dirs_exist_ok=True)
221
222        # Save analysis results if provided
223        if analysis_data is not None and analysis_method is not None:
224            csv_path = dataset_dir / f"{analysis_method.lower()}_summary.csv"
225            _save_analysis_csv(analysis_data, csv_path)
226
227        # Register the hand-annotated alignment entry if a path was provided
228        if hand_alignments_path:
229            from voxkit.storage.alignments import create_hand_alignment
230
231            ok, result = create_hand_alignment(metadata["id"], tg_path=hand_alignments_path)
232            if not ok:
233                raise RuntimeError(f"Failed to register hand alignment: {result}")
234
235        return True, metadata
236
237    except Exception as e:
238        # Clean up on failure
239        dataset_dir = _get_datasets_root() / now
240        if dataset_dir.exists():
241            shutil.rmtree(dataset_dir, ignore_errors=False)
242
243        logger.exception("Error during dataset creation")
244        return False, f"Failed to create dataset metadata: {str(e)}"
245
246
247def _save_analysis_csv(data: list[dict[str, Any]], path: Path) -> None:
248    """Save analysis data to a CSV file.
249
250    Args:
251        data: List of dictionaries where each dictionary represents a row
252        path: Output path for the CSV file
253
254    Raises:
255        ValueError: If data is empty
256    """
257    if not data:
258        raise ValueError("No data to write to CSV.")
259
260    fieldnames = data[0].keys()
261    with open(path, "w", newline="", encoding="utf-8") as csvfile:
262        writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
263        writer.writeheader()
264        for row in data:
265            writer.writerow(row)
266
267
268def get_dataset_metadata(dataset_id: str) -> DatasetMetadata | None:
269    """Get the metadata for a specific dataset.
270
271    Retrieves the dataset metadata from the voxkit_dataset.json file in the
272    dataset's directory.
273
274    Args:
275        dataset_id: ID of the dataset to retrieve
276
277    Returns:
278        Dataset metadata dictionary or None if not found
279
280    Raises:
281        Exception: If metadata file exists but cannot be read or parsed
282    """
283    try:
284        dataset_dir = _get_datasets_root() / dataset_id
285        metadata = _get_dataset_metadata(dataset_dir)
286        if metadata is None:
287            raise FileNotFoundError(f"Metadata for dataset '{dataset_id}' not found.")
288        return metadata
289
290    except Exception:
291        logger.exception("Error retrieving dataset metadata")
292        return None
293
294
295def list_datasets_metadata() -> List[DatasetMetadata]:
296    """List all existing datasets.
297
298    Scans the datasets root directory and collects metadata from all subdirectories
299    containing valid voxkit_dataset.json files.
300
301    Returns:
302        List of dataset metadata dictionaries (empty list if none found)
303
304    Notes:
305        - Silently skips directories without metadata files
306        - Returns empty list on error
307        - Does not guarantee ordering
308    """
309    datasets = []
310    datasets_root = _get_datasets_root()
311
312    try:
313        for entry in os.scandir(datasets_root):
314            if entry.is_dir():
315                metadata_path = os.path.join(entry.path, "voxkit_dataset.json")
316                if os.path.exists(metadata_path):
317                    with open(metadata_path, "r", encoding="utf-8") as f:
318                        metadata = json.load(f)
319                        datasets.append(metadata)
320        return datasets
321
322    except Exception:
323        logger.exception("Error listing datasets")
324        return []
325
326
327def update_dataset_metadata(
328    dataset_id: str,
329    updates: dict,
330) -> Tuple[bool, str]:
331    """Update the metadata for a specific dataset.
332
333    Updates specific fields in the dataset metadata file. Only updates fields that
334    are present in the updates dictionary and not None. Supported fields:
335    description, cached, anonymize, transcribed.
336
337    Args:
338        dataset_id: ID of the dataset to update
339        updates: Dictionary of metadata fields to update (only non-None values are applied)
340
341    Returns:
342        Tuple of (True, success_message) on success or (False, error_message) on failure
343
344    Raises:
345        FileNotFoundError: If the dataset is not found
346        Exception: If metadata file cannot be written
347    """
348    try:
349        metadata = get_dataset_metadata(dataset_id)
350
351        if not metadata:
352            return False, f"Dataset {dataset_id} not found"
353
354        for field in ("description", "cached", "anonymize", "transcribed"):
355            if field in updates and updates[field] is not None:
356                metadata[field] = updates[field]
357
358        # Save the updated metadata
359        metadata_path = _get_datasets_root() / dataset_id / "voxkit_dataset.json"
360        with open(metadata_path, "w", encoding="utf-8") as f:
361            json.dump(metadata, f, indent=2)
362
363        return True, "Dataset metadata updated successfully"
364
365    except KeyError as e:
366        return False, f"Invalid metadata key: {str(e)}"
367    except FileNotFoundError as e:
368        return False, str(e)
369    except Exception as e:
370        return False, f"Failed to update dataset metadata: {str(e)}"
371
372
373def delete_dataset(dataset_id: str) -> Tuple[bool, str]:
374    """Delete a registered dataset.
375
376    Permanently removes the dataset directory and all its contents, including
377    metadata, alignments, and cached data.
378
379    Args:
380        dataset_id: ID of the dataset to delete
381
382    Returns:
383        Tuple of (True, success_message) on success or (False, error_message) on failure
384
385    Raises:
386        Exception: If the directory cannot be removed
387
388    Notes:
389        - This operation is irreversible
390        - Removes the entire dataset directory tree
391        - Validates that dataset_id is not empty before proceeding
392    """
393    if not dataset_id:
394        return False, "Dataset ID cannot be empty."
395
396    dataset_path = _get_datasets_root() / dataset_id
397
398    if dataset_path is None:
399        return False, f"Dataset '{dataset_id}' not found"
400
401    if not dataset_path.exists():
402        return False, f"Dataset '{dataset_id}' not found"
403
404    try:
405        shutil.rmtree(dataset_path)
406        return True, f"Dataset '{dataset_id}' metadata deleted successfully"
407
408    except Exception as e:
409        return False, f"Failed to delete dataset: {str(e)}"
410
411
412def export_dataset(dataset_id: str, output_root: Path) -> Tuple[bool, str]:
413    """Export an existing dataset to a specified output path.
414
415    Copies the entire dataset directory (including metadata, alignments, and cache)
416    to the specified output location. The exported directory is named using the
417    pattern: {dataset_name}_{dataset_id}
418
419    Args:
420        dataset_id: Identifier of the dataset to export
421        output_root: Path to the output directory where the dataset will be copied
422
423    Returns:
424        Tuple of (True, success_message) on success or (False, error_message) on failure
425
426    Raises:
427        FileExistsError: If destination path already exists
428    """
429
430    if not output_root.exists():
431        return False, f"Output path '{output_root}' does not exist."
432    else:
433        dataset_path = _get_datasets_root() / dataset_id
434
435        if not dataset_path.exists():
436            return False, f"Dataset '{dataset_id}' not found."
437
438        dataset_meta = get_dataset_metadata(dataset_id)
439        if not dataset_meta:
440            return False, f"Metadata for dataset '{dataset_id}' not found."
441
442        dest_path = output_root / (dataset_meta["name"] + "_" + dataset_id)
443        try:
444            shutil.copytree(dataset_path, dest_path, dirs_exist_ok=False)
445            return True, f"Dataset '{dataset_id}' exported successfully to '{dest_path}'."
446        except Exception as e:
447            return False, f"Failed to export dataset: {str(e)}"
448
449
450def _rewrite_imported_alignments(new_dataset_path: Path) -> None:
451    """Rewrite alignment metadata paths after importing a dataset to a new location.
452
453    When a dataset is imported, its directory is copied to a new location under a
454    new dataset id. Any ``local`` alignment has a ``tg_path`` that lives inside
455    the dataset directory and still references the source location. For each such
456    alignment, rewrite ``tg_path`` to ``<new_dataset>/alignments/<alignment_id>/
457    textgrids``. Non-local alignments (``local == False``) store TextGrids at the
458    dataset's ``original_path``, which is unchanged by import, so they are left
459    alone.
460    """
461    alignments_dir = new_dataset_path / ALIGNMENTS_ROOT
462    if not alignments_dir.is_dir():
463        return
464
465    for alignment_dir in alignments_dir.iterdir():
466        if not alignment_dir.is_dir():
467            continue
468        metadata_file = alignment_dir / "voxkit_alignment.json"
469        if not metadata_file.exists():
470            continue
471        try:
472            with open(metadata_file, "r", encoding="utf-8") as f:
473                alignment_metadata = json.load(f)
474        except (OSError, json.JSONDecodeError) as e:
475            logger.warning("Skipping alignment metadata rewrite for '%s': %s", metadata_file, e)
476            continue
477
478        if not alignment_metadata.get("local"):
479            continue
480
481        alignment_metadata["tg_path"] = str(alignment_dir / "textgrids")
482        try:
483            with open(metadata_file, "w", encoding="utf-8") as f:
484                json.dump(alignment_metadata, f, indent=4)
485        except OSError:
486            logger.exception("Failed to rewrite alignment metadata '%s'", metadata_file)
487
488
489def import_dataset(dataset_path: Path) -> Tuple[bool, str]:
490    """Import an existing dataset into VoxKit storage.
491
492    Imports a previously exported dataset or a dataset with valid VoxKit metadata.
493    Generates a new ID, updates metadata with new registration date, validates
494    cached datasets, and verifies original path accessibility for non-cached datasets.
495
496    Args:
497        dataset_path: Path to the dataset to import (must contain voxkit_dataset.json)
498
499    Returns:
500        Tuple of (True, success_message) on success or (False, error_message) on failure
501
502    Raises:
503        Exception: If dataset structure is invalid or copy operations fail
504
505    Notes:
506        - For cached datasets, validates the cache directory structure
507        - For non-cached datasets, verifies the original_path is accessible
508        - Automatically cleans up on failure
509    """
510    # Validate dataset structure
511
512    if not isinstance(dataset_path, Path):
513        dataset_path = Path(dataset_path)
514    if not dataset_path.exists():
515        return False, f"Dataset path '{dataset_path}' does not exist."
516    if not dataset_path.is_dir():
517        return False, f"Dataset path '{dataset_path}' is not a directory."
518    dataset_metadata_typed = _get_dataset_metadata(dataset_path)
519    if dataset_metadata_typed is None:
520        return False, "Dataset metadata file not found in the provided dataset path."
521
522    transcribed_flag: bool = bool(dataset_metadata_typed.get("transcribed", True))
523    valid, valid_msg = validate_dataset(dataset_path / "cache", transcribed=transcribed_flag)
524
525    now = generate_unique_id()
526
527    dataset_dest = _get_datasets_root() / now
528    try:
529        # Change metadata accordingly
530        dataset_metadata: dict[str, Any] = dict(dataset_metadata_typed)  # Make a copy to modify
531        dataset_metadata["id"] = now
532        humannow = readable_from_unique_id(now)
533        dataset_metadata["registration_date"] = humannow
534
535        # Check cache consistency
536        if not dataset_metadata["cached"]:
537            original_location_exists = Path(dataset_metadata["original_path"]).exists()
538            if not original_location_exists:
539                return (
540                    False,
541                    f"Original dataset path {dataset_metadata['original_path']} "
542                    "does not exist; cannot import non-cached dataset.",
543                )
544
545        # Validate dataset
546        elif not valid:
547            return False, f"Dataset validation failed: {valid_msg}"
548
549        metadata_path = dataset_dest / "voxkit_dataset.json"
550
551        if not dataset_dest.exists():
552            dataset_dest.mkdir(parents=False, exist_ok=False)
553
554        shutil.copytree(dataset_path, dataset_dest, dirs_exist_ok=True)
555
556        with open(metadata_path, "w", encoding="utf-8") as f:
557            json.dump(dataset_metadata, f, indent=2)
558
559        _rewrite_imported_alignments(dataset_dest)
560
561        return True, "Dataset imported successfully."
562
563    except Exception as e:
564        # Cleanup on failure
565        if dataset_dest.exists():
566            shutil.rmtree(dataset_dest, ignore_errors=True)
567        logger.exception("Error during dataset import")
568        return False, f"Failed to import dataset: {str(e)}"
569
570
571def validate_dataset(dataset_path: Path, transcribed: bool = True) -> Tuple[bool, str]:
572    """Validate if a dataset follows the expected organization pattern.
573
574    Validation checks:
575    - Dataset path exists and is a directory
576    - Dataset is not empty
577    - Contains speaker subdirectories (not files at root level)
578    - Each speaker directory is not empty
579    - Each speaker directory contains audio files (.wav, .flac, .mp3, .ogg, .m4a)
580    - If transcribed=True: each audio file has a matching .lab file in the same directory
581
582    Expected structure:
583
584        dataset_path/
585        ├── speaker_001/
586        │   ├── audio_001.wav
587        │   ├── audio_001.lab  (only required when transcribed=True)
588        │   ├── audio_002.wav
589        │   └── audio_002.lab
590        └── speaker_002/
591            ├── audio_001.wav
592            └── audio_001.lab
593
594    Args:
595        dataset_path: Path to dataset root directory
596        transcribed: Whether to require a .lab file for every audio file
597
598    Returns:
599        Tuple of (True, validation_message) if valid or (False, error_description) if invalid
600    """
601    if not isinstance(dataset_path, Path):
602        dataset_path = Path(dataset_path)
603    if not dataset_path.exists():
604        return False, f"Dataset path '{dataset_path}' does not exist."
605    if not dataset_path.is_dir():
606        return False, f"Dataset path '{dataset_path}' is not a directory."
607    if not os.listdir(dataset_path):
608        return False, f"Dataset path '{dataset_path}' is empty."
609    for subdir in os.listdir(dataset_path):
610        if subdir.startswith("."):
611            continue  # Skip hidden files/directories
612        subdir_path = os.path.join(dataset_path, subdir)
613        if not os.path.isdir(subdir_path):
614            return (
615                False,
616                f"Expected speaker directories in dataset path '{dataset_path}', "
617                f"found file '{subdir_path}'.",
618            )
619        if not os.listdir(subdir_path):
620            return False, f"Speaker directory '{subdir_path}' is empty."
621
622    speaker_dirs = [
623        d for d in os.listdir(dataset_path) if os.path.isdir(os.path.join(dataset_path, d))
624    ]
625
626    if not speaker_dirs:
627        return False, "No speaker directories found in the dataset path."
628
629    for speaker in speaker_dirs:
630        speaker_path = os.path.join(dataset_path, speaker)
631        audio_files = [
632            f for f in os.listdir(speaker_path) if f.endswith(tuple(SUPERSET_AUDIO_EXTENSIONS))
633        ]
634
635        if not audio_files:
636            return False, f"No audio files found in speaker directory '{speaker_path}'."
637
638        lab_files = [f for f in os.listdir(speaker_path) if f.endswith(".lab")]
639
640        if transcribed:
641            audio_stems = {Path(f).stem for f in audio_files}
642            lab_stems = {Path(f).stem for f in lab_files}
643            missing = audio_stems - lab_stems
644            if missing:
645                return (
646                    False,
647                    f"Missing .lab files for audio files in speaker directory "
648                    f"'{speaker_path}': {', '.join(sorted(missing))}.",
649                )
650        else:
651            if lab_files:
652                return (
653                    False,
654                    f"Dataset is marked as not transcribed but .lab files were found in "
655                    f"'{speaker_path}'. Set 'Transcribed' to true or remove the .lab files.",
656                )
657
658    return True, "Dataset is valid."