| 497 | |
| 498 | |
| 499 | class _StateContextManager: |
| 500 | def __init__( |
| 501 | self, |
| 502 | name, |
| 503 | help, |
| 504 | update_thread_local_hook, |
| 505 | validate_new_val_hook: Optional[Callable[[Any], None]] = None, |
| 506 | extra_description: str = "", |
| 507 | default_value: Any = no_default, |
| 508 | ): |
| 509 | self._name = name |
| 510 | self.__name__ = name[4:] if name.startswith("xla_") else name |
| 511 | self.__doc__ = ( |
| 512 | f"Context manager for `{name}` config option" |
| 513 | f"{extra_description}.\n\n{help}" |
| 514 | ) |
| 515 | self._update_thread_local_hook = update_thread_local_hook |
| 516 | self._validate_new_val_hook = validate_new_val_hook |
| 517 | self._default_value = default_value |
| 518 | |
| 519 | @contextlib.contextmanager |
| 520 | def __call__(self, new_val: Any = no_default): |
| 521 | if new_val is no_default: |
| 522 | if self._default_value is not no_default: |
| 523 | new_val = self._default_value # default_value provided to constructor |
| 524 | else: |
| 525 | # no default_value provided to constructor and no value provided as an |
| 526 | # argument, so we raise an error |
| 527 | raise TypeError( |
| 528 | f"Context manager for {self.__name__} config option " |
| 529 | "requires an argument representing the new value for " |
| 530 | "the config option." |
| 531 | ) |
| 532 | if self._validate_new_val_hook: |
| 533 | self._validate_new_val_hook(new_val) |
| 534 | prev_val = getattr(_thread_local_state, self._name, unset) |
| 535 | setattr(_thread_local_state, self._name, new_val) |
| 536 | if self._update_thread_local_hook: |
| 537 | self._update_thread_local_hook(new_val) |
| 538 | try: |
| 539 | yield |
| 540 | finally: |
| 541 | if prev_val is unset: |
| 542 | delattr(_thread_local_state, self._name) |
| 543 | if self._update_thread_local_hook: |
| 544 | self._update_thread_local_hook(None) |
| 545 | else: |
| 546 | setattr(_thread_local_state, self._name, prev_val) |
| 547 | if self._update_thread_local_hook: |
| 548 | self._update_thread_local_hook(prev_val) |
| 549 | |
| 550 | def _add_hooks(self, update_global_hook, update_thread_local_hook): |
| 551 | """Private method that adds hooks to an existing context-manager. |
| 552 | |
| 553 | Used to avoid cyclic import dependencies.""" |
| 554 | self._update_thread_local_hook = update_thread_local_hook |
| 555 | config._update_hooks[self._name] = update_global_hook |
| 556 | update_global_hook(config._read(self._name)) |
no outgoing calls
no test coverage detected