From aa6f229e6bd8d34a0b40c2c08988798a5e5a8efd Mon Sep 17 00:00:00 2001
From: Florian Bruhin <git@the-compiler.org>
Date: Tue, 13 Jun 2017 13:05:24 +0200
Subject: [PATCH] Add utils.yaml_{load,dump}

---
 qutebrowser/config/configdata.py     |  4 +---
 qutebrowser/misc/sessions.py         |  9 ++-------
 qutebrowser/utils/utils.py           | 16 ++++++++++++++++
 tests/unit/config/test_configdata.py |  4 ++++
 tests/unit/utils/test_utils.py       |  8 ++++++++
 5 files changed, 31 insertions(+), 10 deletions(-)

diff --git a/qutebrowser/config/configdata.py b/qutebrowser/config/configdata.py
index bda605642..eea7ec01d 100644
--- a/qutebrowser/config/configdata.py
+++ b/qutebrowser/config/configdata.py
@@ -33,8 +33,6 @@ import sys
 import re
 import collections
 
-import yaml
-
 from qutebrowser.config import configtypes, sections
 from qutebrowser.utils import usertypes, qtutils, utils
 
@@ -648,7 +646,7 @@ def _read_yaml(yaml_data):
         A dict mapping option names to Option elements.
     """
     parsed = {}
-    data = yaml.load(yaml_data)
+    data = utils.yaml_load(yaml_data)
 
     keys = {'type', 'default', 'desc', 'backend'}
 
diff --git a/qutebrowser/misc/sessions.py b/qutebrowser/misc/sessions.py
index 3bbee29f9..c9abfc830 100644
--- a/qutebrowser/misc/sessions.py
+++ b/qutebrowser/misc/sessions.py
@@ -26,10 +26,6 @@ import sip
 from PyQt5.QtCore import pyqtSignal, QUrl, QObject, QPoint, QTimer
 from PyQt5.QtWidgets import QApplication
 import yaml
-try:
-    from yaml import CSafeLoader as YamlLoader, CSafeDumper as YamlDumper
-except ImportError:  # pragma: no cover
-    from yaml import SafeLoader as YamlLoader, SafeDumper as YamlDumper
 
 from qutebrowser.utils import (standarddir, objreg, qtutils, log, usertypes,
                                message, utils)
@@ -298,8 +294,7 @@ class SessionManager(QObject):
         log.sessions.vdebug("Saving data: {}".format(data))
         try:
             with qtutils.savefile_open(path) as f:
-                yaml.dump(data, f, Dumper=YamlDumper, default_flow_style=False,
-                          encoding='utf-8', allow_unicode=True)
+                utils.yaml_dump(data, f)
         except (OSError, UnicodeEncodeError, yaml.YAMLError) as e:
             raise SessionError(e)
         else:
@@ -387,7 +382,7 @@ class SessionManager(QObject):
         path = self._get_session_path(name, check_exists=True)
         try:
             with open(path, encoding='utf-8') as f:
-                data = yaml.load(f, Loader=YamlLoader)
+                data = utils.yaml_load(f, Loader=YamlLoader)
         except (OSError, UnicodeDecodeError, yaml.YAMLError) as e:
             raise SessionError(e)
 
diff --git a/qutebrowser/utils/utils.py b/qutebrowser/utils/utils.py
index 114be6a52..fbafa5353 100644
--- a/qutebrowser/utils/utils.py
+++ b/qutebrowser/utils/utils.py
@@ -36,6 +36,11 @@ from PyQt5.QtCore import Qt, QUrl
 from PyQt5.QtGui import QKeySequence, QColor, QClipboard, QDesktopServices
 from PyQt5.QtWidgets import QApplication
 import pkg_resources
+import yaml
+try:
+    from yaml import CSafeLoader as YamlLoader, CSafeDumper as YamlDumper
+except ImportError:  # pragma: no cover
+    from yaml import SafeLoader as YamlLoader, SafeDumper as YamlDumper
 
 import qutebrowser
 from qutebrowser.utils import qtutils, log
@@ -876,3 +881,14 @@ def expand_windows_drive(path):
         return path + "\\"
     else:
         return path
+
+
+def yaml_load(f):
+    """Wrapper over yaml.load using the C loader if possible."""
+    return yaml.load(f, Loader=YamlLoader)
+
+
+def yaml_dump(data, f=None):
+    """Wrapper over yaml.dump using the C dumper if possible."""
+    return yaml.dump(data, f, Dumper=YamlDumper, default_flow_style=False,
+                     encoding='utf-8', allow_unicode=True)
diff --git a/tests/unit/config/test_configdata.py b/tests/unit/config/test_configdata.py
index 81b680f08..87231f82f 100644
--- a/tests/unit/config/test_configdata.py
+++ b/tests/unit/config/test_configdata.py
@@ -34,6 +34,10 @@ def test_init():
     assert 'ignore_case' in configdata.DATA
 
 
+def test_init_benchmark(benchmark):
+    benchmark(configdata.init)
+
+
 class TestReadYaml:
 
 
diff --git a/tests/unit/utils/test_utils.py b/tests/unit/utils/test_utils.py
index 4b03dcf8c..aff5b49b4 100644
--- a/tests/unit/utils/test_utils.py
+++ b/tests/unit/utils/test_utils.py
@@ -936,3 +936,11 @@ class TestOpenFile:
 ])
 def test_expand_windows_drive(path, expected):
     assert utils.expand_windows_drive(path) == expected
+
+
+def test_yaml_load():
+    assert utils.yaml_load("[1, 2]") == [1, 2]
+
+
+def test_yaml_dump():
+    assert utils.yaml_dump([1, 2]) == b'- 1\n- 2\n'