54 lines
1.3 KiB
Python
54 lines
1.3 KiB
Python
# Use CSafeFile if available
|
|
from typing import TYPE_CHECKING
|
|
from typing import Any
|
|
from typing import Dict
|
|
from typing import Iterable
|
|
from typing import Tuple
|
|
|
|
if TYPE_CHECKING:
|
|
from yaml import SafeLoader
|
|
else:
|
|
try:
|
|
from yaml import CSafeLoader as SafeLoader
|
|
except ImportError:
|
|
from yaml import SafeLoader
|
|
|
|
|
|
__all__ = [
|
|
"SafeLoader",
|
|
]
|
|
|
|
|
|
class LimitedSafeLoader(type):
|
|
"""Meta YAML loader that skips the resolution of the specified YAML tags."""
|
|
|
|
def __new__(
|
|
cls,
|
|
name: str,
|
|
bases: Tuple[type, ...],
|
|
namespace: Dict[str, Any],
|
|
exclude_resolvers: Iterable[str],
|
|
) -> "LimitedSafeLoader":
|
|
exclude_resolvers = set(exclude_resolvers)
|
|
implicit_resolvers = {
|
|
key: [
|
|
(tag, regex)
|
|
for tag, regex in mappings
|
|
if tag not in exclude_resolvers
|
|
]
|
|
for key, mappings in SafeLoader.yaml_implicit_resolvers.items()
|
|
}
|
|
return super().__new__(
|
|
cls,
|
|
name,
|
|
(SafeLoader, *bases),
|
|
{**namespace, "yaml_implicit_resolvers": implicit_resolvers},
|
|
)
|
|
|
|
|
|
class JsonschemaSafeLoader(
|
|
metaclass=LimitedSafeLoader,
|
|
exclude_resolvers={"tag:yaml.org,2002:timestamp"},
|
|
):
|
|
"""A safe YAML loader that leaves timestamps as strings."""
|