Function _wrap_exception

axlearn/common/file_system.py:24–48

Surfaces any `source_exc` as `target_exc` instead. Users can use multiple contexts to wrapping different `target_exc`, e.g: ``` with ( _wrap_exception(tf_errors.OpError, OpError), _wrap_exception(tf_errors.NotFound, NotFoundError), ): ... ``` Note th

(
    source_exc: Union[type[Exception], tuple[type[Exception], ...]], target_exc: type[Exception]
)

@contextlib.contextmanager
def _wrap_exception(
    source_exc: Union[type[Exception], tuple[type[Exception], ...]], target_exc: type[Exception]
):
    """Surfaces any `source_exc` as `target_exc` instead.

    Users can use multiple contexts to wrapping different `target_exc`, e.g:
    ```
    with (
        _wrap_exception(tf_errors.OpError, OpError),
        _wrap_exception(tf_errors.NotFound, NotFoundError),
    ):
        ...
    ```

    Note that contexts are entered in order and exited in reverse order, so more general exceptions
    should be wrapped first.

    Args:
        source_exc: One or more exceptions to catch.
        target_exc: Exception to raise instead.
    """
    try:
        yield
    except source_exc as e:
        raise target_exc(str(e)) from e


class OpError(Exception):