Handle !relative (and any future constructors) in mkdocs.yml; resolves #199

Author: jimporter

CHANGES.md

@@ -6,6 +6,7 @@
 - When calling `set-default`, you can now pass `--allow-undefined` to set the
   default to a version that doesn't exist yet
 - Add global-level `-q` / `--quiet` option to suppress warning messages
+- Add support for handling `!relative` in `mkdocs.yml`
 
 ### Bug fixes
 - When loading an MkDocs config, mike now runs the `startup` and `shutdown`

mike/mkdocs_utils.py

@@ -12,6 +12,31 @@
 docs_version_var = 'MIKE_DOCS_VERSION'
 
 
+class RoundTrippableTag:
+    def __init__(self, node):
+        self.node = node
+
+    def __repr__(self):
+        return repr(self.node)
+
+    @staticmethod
+    def constructor(loader, suffix, node):
+        return RoundTrippableTag(node)
+
+    @staticmethod
+    def representer(dumper, data):
+        return data.node
+
+
+class RoundTripLoader(yaml.SafeLoader):
+    pass
+
+
+yaml.add_multi_constructor('!', RoundTrippableTag.constructor,
+                           Loader=RoundTripLoader)
+yaml.add_multi_representer(RoundTrippableTag, RoundTrippableTag.representer)
+
+
 def _open_config(config_file=None):
     if config_file is None:
         config_file = ['mkdocs.yml', 'mkdocs.yaml']
@@ -45,7 +70,7 @@ def load_config(config_file=None, **kwargs):
 def inject_plugin(config_file):
     with _open_config(config_file) as f:
         config_file = f.name
-        config = mkdocs.utils.yaml_load(f)
+        config = mkdocs.utils.yaml_load(f, loader=RoundTripLoader)
 
     plugins = config.setdefault('plugins', ['search'])
     for i in plugins:

test/unit/test_mkdocs_utils.py

@@ -86,49 +86,49 @@ def test_nonexist(self):
 
 
 class TestInjectPlugin(unittest.TestCase):
+    def setUp(self):
+        self.out = Stream('mike-mkdocs.yml')
+
     def test_no_plugins(self):
-        out = Stream('mike-mkdocs.yml')
         cfg = '{}'
         with mock.patch('builtins.open',
                         mock_open_files({'mkdocs.yml': cfg})), \
              mock.patch('mike.mkdocs_utils.NamedTemporaryFile',
-                        return_value=out), \
+                        return_value=self.out), \
              mock.patch('os.remove') as mremove:
             with mkdocs_utils.inject_plugin('mkdocs.yml') as f:
-                self.assertEqual(f, out.name)
-                newcfg = yaml.load(out.getvalue(), Loader=yaml.Loader)
+                self.assertEqual(f, self.out.name)
+                newcfg = yaml.safe_load(self.out.getvalue())
             mremove.assert_called_once()
 
         self.assertEqual(newcfg, {'plugins': ['mike', 'search']})
 
     def test_other_plugins(self):
-        out = Stream('mike-mkdocs.yml')
         cfg = 'plugins:\n  - foo\n  - bar:\n      option: true'
         with mock.patch('builtins.open',
                         mock_open_files({'mkdocs.yml': cfg})), \
              mock.patch('mike.mkdocs_utils.NamedTemporaryFile',
-                        return_value=out), \
+                        return_value=self.out), \
              mock.patch('os.remove') as mremove:
             with mkdocs_utils.inject_plugin('mkdocs.yml') as f:
-                self.assertEqual(f, out.name)
-                newcfg = yaml.load(out.getvalue(), Loader=yaml.Loader)
+                self.assertEqual(f, self.out.name)
+                newcfg = yaml.safe_load(self.out.getvalue())
             mremove.assert_called_once()
 
         self.assertEqual(newcfg, {'plugins': [
             'mike', 'foo', {'bar': {'option': True}},
         ]})
 
     def test_other_plugins_dict(self):
-        out = Stream('mike-mkdocs.yml')
         cfg = 'plugins:\n  foo: {}\n  bar:\n    option: true'
         with mock.patch('builtins.open',
                         mock_open_files({'mkdocs.yml': cfg})), \
              mock.patch('mike.mkdocs_utils.NamedTemporaryFile',
-                        return_value=out), \
+                        return_value=self.out), \
              mock.patch('os.remove') as mremove:
             with mkdocs_utils.inject_plugin('mkdocs.yml') as f:
-                self.assertEqual(f, out.name)
-                newcfg = yaml.load(out.getvalue(), Loader=yaml.Loader)
+                self.assertEqual(f, self.out.name)
+                newcfg = yaml.safe_load(self.out.getvalue())
             mremove.assert_called_once()
 
         self.assertEqual(newcfg, {'plugins': {
@@ -140,44 +140,62 @@ def test_other_plugins_dict(self):
         )
 
     def test_mike_plugin(self):
-        out = Stream('mike-mkdocs.yml')
         cfg = 'plugins:\n  - mike'
         with mock.patch('builtins.open',
                         mock_open_files({'mkdocs.yml': cfg})), \
              mock.patch('mike.mkdocs_utils.NamedTemporaryFile',
-                        return_value=out), \
+                        return_value=self.out), \
              mock.patch('os.remove') as mremove:
             with mkdocs_utils.inject_plugin('mkdocs.yml') as f:
                 self.assertEqual(f, 'mkdocs.yml')
-                self.assertEqual(out.getvalue(), '')
+                self.assertEqual(self.out.getvalue(), '')
             mremove.assert_not_called()
 
     def test_mike_plugin_options(self):
-        out = Stream('mike-mkdocs.yml')
         cfg = 'plugins:\n  - mike:\n      option: true'
         with mock.patch('builtins.open',
                         mock_open_files({'mkdocs.yml': cfg})), \
              mock.patch('mike.mkdocs_utils.NamedTemporaryFile',
-                        return_value=out), \
+                        return_value=self.out), \
              mock.patch('os.remove') as mremove:
             with mkdocs_utils.inject_plugin('mkdocs.yml') as f:
                 self.assertEqual(f, 'mkdocs.yml')
-                self.assertEqual(out.getvalue(), '')
+                self.assertEqual(self.out.getvalue(), '')
             mremove.assert_not_called()
 
+    def test_round_trip(self):
+        cfg = ('plugins:\n' +
+               '  - foo:\n      option: !relative $config_dir\n' +
+               '  - bar:\n      option: !ENV variable\n' +
+               '  - baz:\n      option: !ENV [variable, default]'
+               )
+        with mock.patch('builtins.open',
+                        mock_open_files({'mkdocs.yml': cfg})), \
+             mock.patch('mike.mkdocs_utils.NamedTemporaryFile',
+                        return_value=self.out), \
+             mock.patch('os.remove') as mremove:
+            with mkdocs_utils.inject_plugin('mkdocs.yml') as f:
+                self.assertEqual(f, self.out.name)
+            mremove.assert_called_once()
+
+        expected = ('plugins:\n- mike\n' +
+                    "- foo:\n    option: !relative '$config_dir'\n" +
+                    "- bar:\n    option: !ENV 'variable'\n"
+                    '- baz:\n    option: !ENV [variable, default]\n')
+        self.assertEqual(self.out.getvalue(), expected)
+
     def test_inherit(self):
-        out = Stream('mike-mkdocs.yml')
         main_cfg = 'INHERIT: mkdocs-base.yml\nplugins:\n  foo: {}\n'
         base_cfg = 'plugins:\n  bar: {}\n'
         files = {'mkdocs.yml': main_cfg, 'mkdocs-base.yml': base_cfg}
         with mock.patch('builtins.open', mock_open_files(files)), \
              mock.patch('mike.mkdocs_utils.NamedTemporaryFile',
-                        return_value=out), \
+                        return_value=self.out), \
              mock.patch('os.path.exists', return_value=True), \
              mock.patch('os.remove') as mremove:
             with mkdocs_utils.inject_plugin('mkdocs.yml') as f:
                 self.assertEqual(f, 'mike-mkdocs.yml')
-                newcfg = yaml.load(out.getvalue(), Loader=yaml.Loader)
+                newcfg = yaml.safe_load(self.out.getvalue())
             mremove.assert_called_once()
 
         self.assertEqual(newcfg, {'plugins': {