From 3351d5feabff8ee107f4ad6d1f86055843c7dbf1 Mon Sep 17 00:00:00 2001
From: Jim Porter <itsjimporter@gmail.com>
Date: Fri, 3 May 2024 09:40:51 -0700
Subject: [PATCH] Expand environment variables when injecting the mike plugin;
 resolves #217

---
 CHANGES.md                     |  8 ++++++++
 mike/mkdocs_utils.py           |  7 +++++++
 setup.py                       |  2 +-
 test/unit/test_mkdocs_utils.py | 34 +++++++++++++++++++++++++++++-----
 4 files changed, 45 insertions(+), 6 deletions(-)

diff --git a/CHANGES.md b/CHANGES.md
index f5441a4..a0fed18 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -1,5 +1,13 @@
 # Changes
 
+## v2.1.1 (in progress)
+
+### Bug fixes
+- Support using environment variables for `INHERIT` when injecting the `mike`
+  plugin into `mkdocs.yml`
+
+---
+
 ## v2.1.0 (2024-05-01)
 
 ### New features
diff --git a/mike/mkdocs_utils.py b/mike/mkdocs_utils.py
index 5051607..811003d 100644
--- a/mike/mkdocs_utils.py
+++ b/mike/mkdocs_utils.py
@@ -5,6 +5,7 @@
 import re
 import subprocess
 import yaml
+import yaml_env_tag
 from collections.abc import Iterable, Mapping
 from contextlib import contextmanager
 from tempfile import NamedTemporaryFile
@@ -32,6 +33,12 @@ class RoundTripLoader(yaml.Loader):
     pass
 
 
+# We need to expand environment variables in our round trip loader (making it
+# less of a "round trip"), or else `INHERIT: !ENV ...` will fail when injecting
+# the mike plugin. MkDocs really doesn't make this easy on us...
+yaml.add_constructor('!ENV', yaml_env_tag.construct_env_tag,
+                     Loader=RoundTripLoader)
+
 yaml.add_multi_constructor('!', RoundTrippableTag.constructor,
                            Loader=RoundTripLoader)
 yaml.add_multi_representer(RoundTrippableTag, RoundTrippableTag.representer)
diff --git a/setup.py b/setup.py
index 680238b..a779e21 100644
--- a/setup.py
+++ b/setup.py
@@ -88,7 +88,7 @@ def run(self):
 
     install_requires=(['importlib_metadata', 'importlib_resources',
                        'jinja2 >= 2.7', 'mkdocs >= 1.0', 'pyparsing >= 3.0',
-                       'pyyaml >= 5.1', 'verspec']),
+                       'pyyaml >= 5.1', 'pyyaml_env_tag', 'verspec']),
     extras_require={
         'dev': ['coverage', 'flake8 >= 3.0', 'flake8-quotes', 'shtab'],
         'test': ['coverage', 'flake8 >= 3.0', 'flake8-quotes', 'shtab'],
diff --git a/test/unit/test_mkdocs_utils.py b/test/unit/test_mkdocs_utils.py
index 235b6e5..f386b52 100644
--- a/test/unit/test_mkdocs_utils.py
+++ b/test/unit/test_mkdocs_utils.py
@@ -167,21 +167,21 @@ 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]'
-               )
+               '  - 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:
+             mock.patch('os.remove') as mremove, \
+             mock.patch.dict(os.environ, {'variable': 'mock_val'}, clear=True):
             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')
+                    '- bar:\n    option: mock_val\n'
+                    '- baz:\n    option: mock_val\n')
         self.assertEqual(self.out.getvalue(), expected)
 
     def test_python_tag(self):
@@ -224,6 +224,30 @@ def test_inherit(self):
             [('mike', {}), ('bar', {}), ('foo', {})]
         )
 
+    def test_inherit_env(self):
+        main_cfg = 'INHERIT: !ENV base_file\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=self.out), \
+             mock.patch('os.path.exists', return_value=True), \
+             mock.patch('os.remove') as mremove, \
+             mock.patch.dict(os.environ, {'base_file': 'mkdocs-base.yml'},
+                             clear=True):
+            with mkdocs_utils.inject_plugin('mkdocs.yml') as f:
+                self.assertEqual(f, 'mike-mkdocs.yml')
+                newcfg = yaml.safe_load(self.out.getvalue())
+            mremove.assert_called_once()
+
+        self.assertEqual(newcfg, {'plugins': {
+            'mike': {}, 'bar': {}, 'foo': {},
+        }})
+        self.assertEqual(
+            list(newcfg['plugins'].items()),
+            [('mike', {}), ('bar', {}), ('foo', {})]
+        )
+
 
 class TestBuild(unittest.TestCase):
     def test_build(self):