yaml.py 2.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172
  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. """
  18. Use libyaml for YAML dump/load operations where possible.
  19. If libyaml is available we will use it -- it is significantly faster.
  20. This module delegates all other properties to the yaml module, so it can be used as:
  21. .. code-block:: python
  22. import airflow.utils.yaml as yaml
  23. And then be used directly in place of the normal python module.
  24. """
  25. from __future__ import annotations
  26. from typing import TYPE_CHECKING, Any, BinaryIO, TextIO, cast
  27. if TYPE_CHECKING:
  28. from yaml.error import MarkedYAMLError, YAMLError # noqa: F401
  29. def safe_load(stream: bytes | str | BinaryIO | TextIO) -> Any:
  30. """Like yaml.safe_load, but use the C libyaml for speed where we can."""
  31. # delay import until use.
  32. from yaml import load as orig
  33. try:
  34. from yaml import CSafeLoader as SafeLoader
  35. except ImportError:
  36. from yaml import SafeLoader # type: ignore[assignment, no-redef]
  37. return orig(stream, SafeLoader)
  38. def dump(data: Any, **kwargs) -> str:
  39. """Like yaml.safe_dump, but use the C libyaml for speed where we can."""
  40. # delay import until use.
  41. from yaml import dump as orig
  42. try:
  43. from yaml import CSafeDumper as SafeDumper
  44. except ImportError:
  45. from yaml import SafeDumper # type: ignore[assignment, no-redef]
  46. return cast(str, orig(data, Dumper=SafeDumper, **kwargs))
  47. def __getattr__(name):
  48. # Delegate anything else to the yaml module
  49. import yaml
  50. if name == "FullLoader":
  51. # Try to use CFullLoader by default
  52. getattr(yaml, "CFullLoader", yaml.FullLoader)
  53. return getattr(yaml, name)