|
18 | 18 | 4. ATCC catalog (type strains, genome links) |
19 | 19 | """ |
20 | 20 |
|
| 21 | +import argparse |
21 | 22 | import re |
22 | | -import yaml |
23 | | -import duckdb |
| 23 | +import sys |
| 24 | +from collections import defaultdict |
| 25 | +from dataclasses import dataclass, field |
24 | 26 | from pathlib import Path |
25 | 27 | from typing import Dict, List, Optional, Set, Tuple |
26 | | -from dataclasses import dataclass, field |
27 | | -import requests |
28 | | -import time |
29 | | -from collections import defaultdict |
| 28 | + |
| 29 | +import duckdb |
| 30 | +import yaml |
| 31 | + |
| 32 | +# Add src to path for imports |
| 33 | +sys.path.insert(0, str(Path(__file__).parent.parent / "src")) |
| 34 | + |
| 35 | +from communitymech.curate.curation_event import record_curation_event |
| 36 | +from communitymech.validation.write_validated import ( |
| 37 | + ValidationFailedError, |
| 38 | + write_validated_community, |
| 39 | +) |
| 40 | + |
30 | 41 |
|
31 | 42 | # Color codes for output |
32 | 43 | class Colors: |
@@ -343,13 +354,111 @@ def print_summary(self): |
343 | 354 | print(f" Taxa with culture collections: {self.stats['strains_with_collections']}") |
344 | 355 | print(f" Taxa with genome accessions: {self.stats['strains_with_genome']}") |
345 | 356 |
|
| 357 | + def apply_strain_data_to_community( |
| 358 | + self, |
| 359 | + yaml_path: Path, |
| 360 | + strain_data: Dict[str, StrainInfo], |
| 361 | + *, |
| 362 | + overwrite: bool = False, |
| 363 | + ) -> int: |
| 364 | + """Write extracted strain_designation entries back into a community YAML. |
| 365 | +
|
| 366 | + Loads ``yaml_path``, attaches a ``strain_designation`` to each |
| 367 | + matching ``taxonomy[*].taxon_term`` (matched by ``preferred_term``), |
| 368 | + appends a ``CurationEvent``, and writes via |
| 369 | + :func:`write_validated_community` so closed-schema LinkML validation |
| 370 | + gates the disk write. Returns the number of taxa updated. |
| 371 | +
|
| 372 | + Args: |
| 373 | + yaml_path: Community YAML to update. |
| 374 | + strain_data: Mapping ``preferred_term -> StrainInfo`` produced by |
| 375 | + :meth:`extract_strain_from_yaml`. |
| 376 | + overwrite: When False (default), skip taxa that already carry a |
| 377 | + ``strain_designation`` so curator-authored data is preserved. |
| 378 | +
|
| 379 | + Raises: |
| 380 | + ValidationFailedError: re-raised by the caller for visibility; |
| 381 | + callers in a batch loop should ``except`` and continue. |
| 382 | + """ |
| 383 | + with open(yaml_path) as f: |
| 384 | + data = yaml.safe_load(f) |
| 385 | + |
| 386 | + if 'taxonomy' not in data: |
| 387 | + return 0 |
| 388 | + |
| 389 | + updated_taxa = [] |
| 390 | + for taxon_entry in data['taxonomy']: |
| 391 | + taxon_term = taxon_entry.get('taxon_term') or {} |
| 392 | + preferred_term = taxon_term.get('preferred_term', '') |
| 393 | + if preferred_term not in strain_data: |
| 394 | + continue |
| 395 | + |
| 396 | + if 'strain_designation' in taxon_entry and not overwrite: |
| 397 | + continue |
| 398 | + |
| 399 | + snippet = self.generate_yaml_snippet(strain_data[preferred_term]) |
| 400 | + if not snippet: |
| 401 | + continue |
| 402 | + |
| 403 | + taxon_entry['strain_designation'] = snippet |
| 404 | + updated_taxa.append(preferred_term) |
| 405 | + |
| 406 | + if not updated_taxa: |
| 407 | + return 0 |
| 408 | + |
| 409 | + record_curation_event( |
| 410 | + data, |
| 411 | + curator="enhance_strain_data", |
| 412 | + action="ENHANCE_STRAIN_DATA", |
| 413 | + changes=( |
| 414 | + f"Added strain_designation for {len(updated_taxa)} taxa: " |
| 415 | + f"{', '.join(updated_taxa[:5])}" |
| 416 | + + ("..." if len(updated_taxa) > 5 else "") |
| 417 | + ), |
| 418 | + ) |
| 419 | + |
| 420 | + write_validated_community(data, yaml_path) |
| 421 | + return len(updated_taxa) |
| 422 | + |
346 | 423 |
|
347 | 424 | def main(): |
| 425 | + parser = argparse.ArgumentParser( |
| 426 | + description=( |
| 427 | + "Phase 2: extract strain designations and (optionally) apply " |
| 428 | + "them to community YAMLs" |
| 429 | + ) |
| 430 | + ) |
| 431 | + parser.add_argument( |
| 432 | + '--apply', |
| 433 | + action='store_true', |
| 434 | + help=( |
| 435 | + "Write extracted strain_designation entries back into " |
| 436 | + "kb/communities/*.yaml via write_validated_community(). " |
| 437 | + "Without this flag the script only emits the report + snippets " |
| 438 | + "files for human review (the historical default)." |
| 439 | + ), |
| 440 | + ) |
| 441 | + parser.add_argument( |
| 442 | + '--overwrite', |
| 443 | + action='store_true', |
| 444 | + help=( |
| 445 | + "With --apply, replace existing strain_designation entries. " |
| 446 | + "Default behavior preserves curator-authored values." |
| 447 | + ), |
| 448 | + ) |
| 449 | + parser.add_argument( |
| 450 | + '--kb-dir', |
| 451 | + type=Path, |
| 452 | + default=Path('kb/communities'), |
| 453 | + help="Path to community YAML directory (default: kb/communities)", |
| 454 | + ) |
| 455 | + args = parser.parse_args() |
| 456 | + |
348 | 457 | print(f"{Colors.BOLD}{Colors.CYAN}Phase 2: Data Enhancement - Strain Resolution{Colors.RESET}") |
349 | 458 | print(f"{Colors.CYAN}Strategy: Literature → kg-microbe → APIs{Colors.RESET}\n") |
350 | 459 |
|
351 | 460 | # Paths |
352 | | - kb_dir = Path('/Users/marcin/Documents/VIMSS/ontology/KG-Hub/KG-Microbe/CommunityMech/CommunityMech/kb/communities') |
| 461 | + kb_dir = args.kb_dir |
353 | 462 | kgm_db = Path('kgm_taxonomy.duckdb') |
354 | 463 | output_dir = Path('.') |
355 | 464 |
|
@@ -420,14 +529,56 @@ def main(): |
420 | 529 |
|
421 | 530 | print(f"{Colors.GREEN}✓{Colors.RESET} Written: {snippets_path}") |
422 | 531 |
|
| 532 | + # Apply strain designations to community YAMLs when --apply is set. |
| 533 | + # Without --apply the script keeps its historical "extract + report" |
| 534 | + # behavior and writes nothing to kb/communities/. With --apply each |
| 535 | + # community is loaded, mutated in-memory, gets a CurationEvent appended, |
| 536 | + # and is written via write_validated_community() so closed-schema |
| 537 | + # validation refuses any doc that drifted into an invalid shape. |
| 538 | + if args.apply: |
| 539 | + print(f"\n{Colors.CYAN}Applying strain designations to community YAMLs...{Colors.RESET}") |
| 540 | + applied_total = 0 |
| 541 | + applied_files = 0 |
| 542 | + failed_files = 0 |
| 543 | + for yaml_path, strain_data in sorted(all_strain_data.items()): |
| 544 | + try: |
| 545 | + count = extractor.apply_strain_data_to_community( |
| 546 | + yaml_path, |
| 547 | + strain_data, |
| 548 | + overwrite=args.overwrite, |
| 549 | + ) |
| 550 | + except ValidationFailedError as exc: |
| 551 | + print( |
| 552 | + f" {Colors.RED}✗{Colors.RESET} validation failed for " |
| 553 | + f"{yaml_path.name}: {exc.summary()}", |
| 554 | + file=sys.stderr, |
| 555 | + ) |
| 556 | + failed_files += 1 |
| 557 | + continue |
| 558 | + |
| 559 | + if count > 0: |
| 560 | + applied_total += count |
| 561 | + applied_files += 1 |
| 562 | + print( |
| 563 | + f" {Colors.GREEN}✓{Colors.RESET} {yaml_path.name}: " |
| 564 | + f"applied strain_designation to {count} taxa" |
| 565 | + ) |
| 566 | + |
| 567 | + print( |
| 568 | + f"\n{Colors.GREEN}Applied strain_designation to {applied_total} " |
| 569 | + f"taxa across {applied_files} community file(s); " |
| 570 | + f"{failed_files} file(s) failed validation.{Colors.RESET}" |
| 571 | + ) |
| 572 | + |
423 | 573 | # Print summary |
424 | 574 | extractor.print_summary() |
425 | 575 |
|
426 | 576 | print(f"\n{Colors.GREEN}{Colors.BOLD}✓ Phase 2 strain extraction complete!{Colors.RESET}") |
427 | 577 | print(f"\n{Colors.CYAN}Next steps:{Colors.RESET}") |
428 | 578 | print(f" 1. Review {report_path}") |
429 | 579 | print(f" 2. Review {snippets_path}") |
430 | | - print(f" 3. Apply strain designations to YAML files (Phase 2B)") |
| 580 | + if not args.apply: |
| 581 | + print(f" 3. Apply strain designations to YAML files: re-run with --apply") |
431 | 582 | print(f" 4. Query BacDive/ATCC APIs for additional metadata (Phase 2C)") |
432 | 583 |
|
433 | 584 | if __name__ == '__main__': |
|
0 commit comments