Always use yaml SafeLoader

mattrobenolt · mattrobenolt · commit 898cf8131a72 · 2020-01-09T14:11:16.000-08:00
diff --git a/ua_parser/user_agent_parser_test.py b/ua_parser/user_agent_parser_test.py
@@ -32,6 +32,12 @@
 import unittest
 import yaml
 
+try:
+    # Try and use libyaml bindings if available since faster
+    from yaml import CSafeLoader as SafeLoader
+except ImportError:
+    from yaml import SafeLoader
+
 from ua_parser import user_agent_parser
 
 TEST_RESOURCES_DIR = os.path.join(
@@ -113,7 +119,7 @@ def makePGTSComparisonYAML(self):
         print >> outfile, "test_cases:"
 
         yamlFile = open(os.path.join(TEST_RESOURCES_DIR, "pgts_browser_list.yaml"))
-        yamlContents = yaml.load(yamlFile)
+        yamlContents = yaml.load(yamlFile, Loader=SafeLoader)
         yamlFile.close()
 
         for test_case in yamlContents["test_cases"]:
@@ -140,7 +146,7 @@ def makePGTSComparisonYAML(self):
     # Run a set of test cases from a YAML file
     def runUserAgentTestsFromYAML(self, file_name):
         yamlFile = open(os.path.join(TEST_RESOURCES_DIR, file_name))
-        yamlContents = yaml.load(yamlFile)
+        yamlContents = yaml.load(yamlFile, Loader=SafeLoader)
         yamlFile.close()
 
         for test_case in yamlContents["test_cases"]:
@@ -178,7 +184,7 @@ def runUserAgentTestsFromYAML(self, file_name):
 
     def runOSTestsFromYAML(self, file_name):
         yamlFile = open(os.path.join(TEST_RESOURCES_DIR, file_name))
-        yamlContents = yaml.load(yamlFile)
+        yamlContents = yaml.load(yamlFile, Loader=SafeLoader)
         yamlFile.close()
 
         for test_case in yamlContents["test_cases"]:
@@ -218,7 +224,7 @@ def runOSTestsFromYAML(self, file_name):
 
     def runDeviceTestsFromYAML(self, file_name):
         yamlFile = open(os.path.join(TEST_RESOURCES_DIR, file_name))
-        yamlContents = yaml.load(yamlFile)
+        yamlContents = yaml.load(yamlFile, Loader=SafeLoader)
         yamlFile.close()
 
         for test_case in yamlContents["test_cases"]:
