Allow non-default yaml loader and dumper (#10032)

* Allow non-default yaml loader * Fix style
DataDog · Sep 3, 2021 · e01823a · e01823a
1 parent aa58f6f
commit e01823a
Showing 1 changed file with 15 additions and 25 deletions.
diff --git a/datadog_checks_base/datadog_checks/base/ddyaml.py b/datadog_checks_base/datadog_checks/base/ddyaml.py
@@ -2,10 +2,8 @@
 # All rights reserved
 # Licensed under Simplified BSD License (see LICENSE)
 import logging
-from os.path import realpath
 
 import yaml
-from six import string_types
 
 try:
     from yaml import CSafeDumper as yDumper
@@ -40,14 +38,15 @@ def safe_yaml_dump_all(
     tags=None,
 ):
     if Dumper != yDumper:
-        stream_name = get_stream_name(stream)
-        log.debug("Unsafe dumping of YAML has been disabled - using safe dumper instead in %s", stream_name)
+        log.debug(
+            "`%s` YAML dumper is used instead of the default one, please make sure it is safe to do so", Dumper.__name__
+        )
 
     if pyyaml_dump_all:
         return pyyaml_dump_all(
             documents,
             stream,
-            yDumper,
+            Dumper,
             default_style,
             default_flow_style,
             canonical,
@@ -65,7 +64,7 @@ def safe_yaml_dump_all(
     return yaml.dump_all(
         documents,
         stream,
-        yDumper,
+        Dumper,
         default_style,
         default_flow_style,
         canonical,
@@ -83,35 +82,26 @@ def safe_yaml_dump_all(
 
 def safe_yaml_load(stream, Loader=yLoader):
     if Loader != yLoader:
-        stream_name = get_stream_name(stream)
-        log.debug("Unsafe loading of YAML has been disabled - using safe loader instead in %s", stream_name)
+        log.debug(
+            "`%s` YAML loader is used instead of the default one, please make sure it is safe to do so", Loader.__name__
+        )
 
     if pyyaml_load:
-        return pyyaml_load(stream, Loader=yLoader)
+        return pyyaml_load(stream, Loader=Loader)
 
-    return yaml.load(stream, Loader=yLoader)
+    return yaml.load(stream, Loader=Loader)
 
 
 def safe_yaml_load_all(stream, Loader=yLoader):
     if Loader != yLoader:
-        stream_name = get_stream_name(stream)
-        log.debug("Unsafe loading of YAML has been disabled - using safe loader instead in %s", stream_name)
+        log.debug(
+            "`%s` YAML loader is used instead of the default one, please make sure it is safe to do so", Loader.__name__
+        )
 
     if pyyaml_load_all:
-        return pyyaml_load_all(stream, Loader=yLoader)
-
-    return yaml.load_all(stream, Loader=yLoader)
-
+        return pyyaml_load_all(stream, Loader=Loader)
 
-def get_stream_name(stream):
-    """Using the same logic as pyyaml to handle both string types and file types. All file objects do not necessarily
-    have a `name` attribute, in that case we can only say the stream is a file."""
-    if isinstance(stream, string_types):
-        return "<string>"
-    elif hasattr(stream, 'name'):
-        return realpath(stream.name)
-    else:
-        return "<file>"
+    return yaml.load_all(stream, Loader=Loader)
 
 
 def monkey_patch_pyyaml():