On source install C Extensions might have not been built

checks/__init__.py

@@ -16,7 +16,7 @@
 from pprint import pprint
 from collections import defaultdict
 
-from util import LaconicFilter, get_os, get_hostname, get_next_id
+from util import LaconicFilter, get_os, get_hostname, get_next_id, yLoader
 from config import get_confd_path
 from checks import check_status
 
@@ -585,7 +585,7 @@ def from_yaml(cls, path_to_yaml=None, agentConfig=None, yaml_text=None, check_na
             yaml_text = f.read()
             f.close()
 
-        config = yaml.load(yaml_text, Loader=yaml.CLoader)
+        config = yaml.load(yaml_text, Loader=yLoader)
         check = cls(check_name, config.get('init_config') or {}, agentConfig or {})
 
         return check, config.get('instances', [])

config.py

@@ -20,7 +20,7 @@
 
 # project
 
-from util import get_os, Platform
+from util import get_os, Platform, yLoader
 from jmxfetch import JMXFetch, JMX_COLLECT_COMMAND
 from migration import migrate_old_style_configuration
 
@@ -656,7 +656,7 @@ def check_yaml(conf_path):
     f = open(conf_path)
     check_name = os.path.basename(conf_path).split('.')[0]
     try:
-        check_config = yaml.load(f.read(), Loader=yaml.CLoader)
+        check_config = yaml.load(f.read(), Loader=yLoader)
         assert 'init_config' in check_config, "No 'init_config' section found"
         assert 'instances' in check_config, "No 'instances' section found"
 

jmxfetch.py

@@ -8,7 +8,7 @@
 import time
 
 # datadog
-from util import PidFile, get_os
+from util import PidFile, get_os, yLoader, yDumper
 
 # 3rd party
 import yaml
@@ -86,7 +86,7 @@ def write_status_file(cls, invalid_checks):
             'invalid_checks': invalid_checks
         }
         stream = file(os.path.join(tempfile.gettempdir(), PYTHON_JMX_STATUS_FILE), 'w')
-        yaml.dump(data, stream)
+        yaml.dump(data, stream, Dumper=yDumper)
         stream.close()
 
     @classmethod
@@ -125,7 +125,7 @@ def should_run(cls, confd_path, checks_list):
             if os.path.exists(conf):
                 f = open(conf)
                 try:
-                    check_config = yaml.load(f.read(), Loader=yaml.CLoader)
+                    check_config = yaml.load(f.read(), Loader=yLoader)
                     assert check_config is not None
                     f.close()
                 except Exception:

migration.py

@@ -13,6 +13,9 @@
 import logging
 import string
 
+# project
+from util import yDumper
+
 # 3rd party
 import yaml
 
@@ -294,7 +297,7 @@ def _write_conf(check_name, config, confd_dir):
         raise NoConfigToMigrateException()
 
     try:
-        yaml_config = yaml.dump(config, Dumper=yaml.CDumper, default_flow_style=False)
+        yaml_config = yaml.dump(config, Dumper=yDumper, default_flow_style=False)
     except Exception, e:
         log.exception("Couldn't create yaml from config: %s" % config)
         return

util.py

@@ -18,6 +18,16 @@
 # Tornado
 from tornado import ioloop
 
+# yaml
+import yaml
+try:
+    from yaml import CLoader as yLoader
+    from yaml import CDumper as yDumper
+except ImportError:
+    # On source install C Extensions might have not been built
+    from yaml import Loader as yLoader
+    from yaml import Dumper as yDumper
+
 
 
 VALID_HOSTNAME_RFC_1123_PATTERN = re.compile(r"^(([a-zA-Z0-9]|[a-zA-Z0-9][a-zA-Z0-9\-]*[a-zA-Z0-9])\.)*([A-Za-z0-9]|[A-Za-z0-9][A-Za-z0-9\-]*[A-Za-z0-9])$")