Use C-accelerated safe_load for YAML

paasta_tools/api/api.py

@@ -24,12 +24,12 @@
 import manhole
 import requests_cache
 import service_configuration_lib
-import yaml
 from pyramid.config import Configurator
 from wsgicors import CORS
 
 import paasta_tools.api
 from paasta_tools import kubernetes_tools
+from paasta_tools import yaml_tools as yaml
 from paasta_tools.api import settings
 from paasta_tools.api.tweens import profiling
 from paasta_tools.api.tweens import request_logger

paasta_tools/cli/cmds/validate.py

@@ -33,7 +33,6 @@
 from typing import Union
 
 import pytz
-import yaml
 from croniter import croniter
 from jsonschema import Draft4Validator
 from jsonschema import exceptions
@@ -45,6 +44,7 @@
 from ruamel.yaml import YAML
 from ruamel.yaml.comments import CommentedMap
 
+from paasta_tools import yaml_tools as yaml
 from paasta_tools.autoscaling.utils import MetricsProviderDict
 from paasta_tools.cli.utils import failure
 from paasta_tools.cli.utils import get_file_contents

paasta_tools/cli/fsm/autosuggest.py

@@ -14,7 +14,7 @@
 import os
 import random
 
-import yaml
+from paasta_tools import yaml_tools as yaml
 
 
 def _get_smartstack_proxy_ports_from_file(root, file):

paasta_tools/contrib/check_orphans.py

@@ -19,10 +19,11 @@
 from typing import Set
 from typing import Tuple
 
-import yaml
 from kazoo.client import KazooClient
 from kazoo.exceptions import NoNodeError
 
+from paasta_tools import yaml_tools as yaml
+
 logger = logging.getLogger("check_orphans")
 
 PREFIX = "/smartstack/global/"

paasta_tools/contrib/render_template.py

@@ -3,7 +3,7 @@
 import os
 import re
 
-import yaml
+from paasta_tools import yaml_tools as yaml
 
 
 def replace(s, values):

paasta_tools/contrib/timeouts_metrics_prom.py

@@ -16,12 +16,12 @@
 """
 import os
 
-import yaml
 from prometheus_client import CollectorRegistry
 from prometheus_client import Gauge
 from prometheus_client import write_to_textfile
 from prometheus_client.metrics import MetricWrapperBase
 
+from paasta_tools import yaml_tools as yaml
 from paasta_tools.utils import DEFAULT_SOA_DIR
 
 PROM_OUTPUT_FILE = f"{DEFAULT_SOA_DIR}/.autotune_timeouts.prom"

paasta_tools/envoy_tools.py

@@ -30,10 +30,10 @@
 from typing import Tuple
 
 import requests
-import yaml
 from kubernetes.client import V1Pod
 from mypy_extensions import TypedDict
 
+from paasta_tools import yaml_tools as yaml
 from paasta_tools.utils import get_user_agent
 
 

paasta_tools/generate_authenticating_services.py

@@ -24,8 +24,7 @@
 from typing import List
 from typing import Set
 
-import yaml
-
+from paasta_tools import yaml_tools as yaml
 from paasta_tools.utils import DEFAULT_SOA_DIR
 from paasta_tools.utils import load_system_paasta_config
 from paasta_tools.utils import write_json_configuration_file

paasta_tools/setup_istio_mesh.py

@@ -32,8 +32,8 @@
 from typing import Set
 
 import kubernetes.client as k8s
-import yaml
 
+from paasta_tools import yaml_tools as yaml
 from paasta_tools.kubernetes_tools import ensure_namespace
 from paasta_tools.kubernetes_tools import KubeClient
 from paasta_tools.kubernetes_tools import limit_size_with_hash
@@ -99,7 +99,7 @@ def load_smartstack_namespaces(soa_dir: str = DEFAULT_SOA_DIR) -> Mapping:
             continue
         try:
             with open(file_path) as f:
-                svc_namespaces = yaml.load(f, Loader=yaml.CSafeLoader)
+                svc_namespaces = yaml.safe_load(f)
                 for (ns, details) in svc_namespaces.items():
                     namespaces[f"{dir}.{ns}"] = details
         except Exception as err:

paasta_tools/setup_kubernetes_cr.py

@@ -28,9 +28,9 @@
 from typing import Optional
 from typing import Sequence
 
-import yaml
 from kubernetes.client.exceptions import ApiException
 
+from paasta_tools import yaml_tools as yaml
 from paasta_tools.cli.utils import LONG_RUNNING_INSTANCE_TYPE_HANDLERS
 from paasta_tools.flink_tools import get_flink_ingress_url_root
 from paasta_tools.kubernetes_tools import create_custom_resource

paasta_tools/tron/client.py

@@ -15,8 +15,8 @@
 from urllib.parse import urljoin
 
 import requests
-import yaml
 
+from paasta_tools import yaml_tools as yaml
 from paasta_tools.utils import get_user_agent
 
 

paasta_tools/tron_tools.py

@@ -25,13 +25,13 @@
 from typing import Tuple
 from typing import Union
 
-import yaml
 from mypy_extensions import TypedDict
 from service_configuration_lib import read_extra_service_information
 from service_configuration_lib import read_yaml_file
 from service_configuration_lib.spark_config import get_total_driver_memory_mb
 from service_configuration_lib.spark_config import SparkConfBuilder
 
+from paasta_tools import yaml_tools as yaml
 from paasta_tools.mesos_tools import mesos_services_running_here
 
 try:

paasta_tools/utils.py

@@ -74,7 +74,6 @@
 import ldap3
 import requests_cache
 import service_configuration_lib
-import yaml
 from docker import APIClient
 from docker.utils import kwargs_from_env
 from kazoo.client import KazooClient
@@ -83,6 +82,7 @@
 from service_configuration_lib import read_service_configuration
 
 import paasta_tools.cli.fsm
+from paasta_tools import yaml_tools as yaml
 
 
 # DO NOT CHANGE SPACER, UNLESS YOU'RE PREPARED TO CHANGE ALL INSTANCES

paasta_tools/yaml_tools.py

@@ -0,0 +1,44 @@
+import sys
+
+import yaml
+
+# try and catch both /opt/venvs/paasta-tools and ~/pg/paasta/.tox/py38-linux as if we're being run as an application,
+# we likely want to fail on a potential slowdown rather than experience a performance regression
+if "paasta" in sys.prefix:
+    from yaml import CSafeLoader as Loader
+    from yaml import CSafeDumper as Dumper
+# but for the vanishingly few instances where folks add us as as dependency, we don't enforce that they use the
+# C-accelerated Loader/Dumper
+else:
+    try:
+        from yaml import CSafeLoader as Loader
+        from yaml import CSafeDumper as Dumper
+    except ImportError:  # pragma: no cover
+        from yaml import SafeLoader as Loader  # type: ignore
+        from yaml import SafeDumper as Dumper  # type: ignore
+
+
+def dump(*args, **kwargs):
+    kwargs["Dumper"] = Dumper
+    return yaml.dump(*args, **kwargs)
+
+
+def dump_all(*args, **kwargs):
+    kwargs["Dumper"] = Dumper
+    return yaml.dump_all(*args, **kwargs)
+
+
+def load(*args, **kwargs):
+    kwargs["Loader"] = Loader
+    return yaml.load(*args, **kwargs)
+
+
+def load_all(*args, **kwargs):
+    kwargs["Loader"] = Loader
+    return yaml.load_all(*args, **kwargs)
+
+
+safe_dump = dump
+safe_dump_all = dump_all
+safe_load = load
+safe_load_all = load_all

tests/test_config_utils.py

@@ -2,9 +2,9 @@
 
 import mock
 import pytest
-import yaml
 
 import paasta_tools.config_utils as config_utils
+from paasta_tools import yaml_tools as yaml
 from paasta_tools.utils import AUTO_SOACONFIG_SUBDIR
 
 

tests/test_firewall_update.py

@@ -15,10 +15,10 @@
 
 import mock
 import pytest
-import yaml
 
 from paasta_tools import firewall
 from paasta_tools import firewall_update
+from paasta_tools import yaml_tools as yaml
 from paasta_tools.utils import TimeoutError
 
 

tests/test_generate_authenticating_services.py

@@ -14,9 +14,9 @@
 from unittest.mock import patch
 
 import pytest
-import yaml
 
 from paasta_tools import generate_authenticating_services
+from paasta_tools import yaml_tools as yaml
 
 
 @pytest.fixture

tests/test_tron_tools.py

@@ -5,10 +5,10 @@
 
 import mock
 import pytest
-import yaml
 
 from paasta_tools import tron_tools
 from paasta_tools import utils
+from paasta_tools import yaml_tools as yaml
 from paasta_tools.secret_tools import SHARED_SECRET_SERVICE
 from paasta_tools.tron_tools import MASTER_NAMESPACE
 from paasta_tools.tron_tools import MESOS_EXECUTOR_NAMES