warnings.py 1.6 KB

12345678910111213141516171819202122232425262728293031323334353637383940
  1. # Licensed to the Apache Software Foundation (ASF) under one
  2. # or more contributor license agreements. See the NOTICE file
  3. # distributed with this work for additional information
  4. # regarding copyright ownership. The ASF licenses this file
  5. # to you under the Apache License, Version 2.0 (the
  6. # "License"); you may not use this file except in compliance
  7. # with the License. You may obtain a copy of the License at
  8. #
  9. # http://www.apache.org/licenses/LICENSE-2.0
  10. #
  11. # Unless required by applicable law or agreed to in writing,
  12. # software distributed under the License is distributed on an
  13. # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
  14. # KIND, either express or implied. See the License for the
  15. # specific language governing permissions and limitations
  16. # under the License.
  17. from __future__ import annotations
  18. import warnings
  19. from collections.abc import Generator
  20. from contextlib import contextmanager
  21. @contextmanager
  22. def capture_with_reraise() -> Generator[list[warnings.WarningMessage], None, None]:
  23. """Capture warnings in context and re-raise it on exit from the context manager."""
  24. captured_warnings = []
  25. try:
  26. with warnings.catch_warnings(record=True) as captured_warnings:
  27. yield captured_warnings
  28. finally:
  29. if captured_warnings:
  30. for cw in captured_warnings:
  31. warnings.warn_explicit(
  32. message=cw.message,
  33. category=cw.category,
  34. filename=cw.filename,
  35. lineno=cw.lineno,
  36. source=cw.source,
  37. )