From 83dea299dd3614d1911e2f4d5bde9148d907b8f0 Mon Sep 17 00:00:00 2001
From: YuviPanda <yuvipanda@gmail.com>
Date: Thu, 14 Sep 2023 12:20:06 -0700
Subject: [PATCH] Support lists & dicts in unlisted_choice kubespawner_override

---
 kubespawner/spawner.py |  7 ++-----
 kubespawner/utils.py   | 43 ++++++++++++++++++++++++++++++++++++++++++
 tests/test_profile.py  | 16 +++++++++++++++-
 3 files changed, 60 insertions(+), 6 deletions(-)

diff --git a/kubespawner/spawner.py b/kubespawner/spawner.py
index fa511e33..67e8fc24 100644
--- a/kubespawner/spawner.py
+++ b/kubespawner/spawner.py
@@ -46,7 +46,7 @@
     make_service,
 )
 from .reflector import ResourceReflector
-from .utils import recursive_update
+from .utils import recursive_format, recursive_update
 
 
 class PodReflector(ResourceReflector):
@@ -3191,10 +3191,7 @@ def _load_profile(self, slug, profile_list):
                     "kubespawner_override", {}
                 )
                 for k, v in option_overrides.items():
-                    # FIXME: This logic restricts unlisted_choice to define
-                    #        kubespawner_override dictionaries where all keys
-                    #        have string values.
-                    option_overrides[k] = v.format(value=unlisted_choice)
+                    option_overrides[k] = recursive_format(v, value=unlisted_choice)
             elif choice:
                 # A pre-defined choice was selected
                 option_overrides = option["choices"][choice].get(
diff --git a/kubespawner/utils.py b/kubespawner/utils.py
index 47db20eb..a9335861 100644
--- a/kubespawner/utils.py
+++ b/kubespawner/utils.py
@@ -3,6 +3,7 @@
 """
 import copy
 import hashlib
+from typing import Union
 
 
 def generate_hashed_slug(slug, limit=63, hash_length=6):
@@ -222,3 +223,45 @@ def recursive_update(target, new):
 
         else:
             target[k] = v
+
+
+class IgnoreMissing(dict):
+    """
+    Dictionary subclass for use with format_map
+
+    Renders missing values as "{key}", so format strings with
+    missing values just get rendered as is.
+
+    Stolen from https://docs.python.org/3/library/stdtypes.html#str.format_map
+    """
+
+    def __missing__(self, key):
+        return f"{{{key}}}"
+
+
+def recursive_format(
+    format_object: Union[str, dict, list], **kwargs
+) -> Union[str, dict, list]:
+    """
+    Recursively format given object with values provided as keyword arguments.
+
+    If the given object (string, list or dict) has items that do not have
+    placeholders for passed in kwargs, no formatting is performed.
+
+    recursive_format("{v}", v=5) -> Returns "5"
+    recrusive_format("{a}") -> Returns "{a}" rather than erroring, as is
+    the behavior of "format"
+    """
+    if isinstance(format_object, str):
+        return format_object.format_map(IgnoreMissing(kwargs))
+    elif isinstance(format_object, list):
+        return [recursive_format(i, **kwargs) for i in format_object]
+    elif isinstance(format_object, dict):
+        return {
+            recursive_format(k, **kwargs): recursive_format(v, **kwargs)
+            for k, v in format_object.items()
+        }
+    else:
+        raise ValueError(
+            f"Object of type {type(format_object)} can not be recursively formatted"
+        )
diff --git a/tests/test_profile.py b/tests/test_profile.py
index f40c9cc5..87a3d40d 100644
--- a/tests/test_profile.py
+++ b/tests/test_profile.py
@@ -256,7 +256,15 @@ async def test_unlisted_choice_non_string_override():
                             'environment': {
                                 'CUSTOM_IMAGE_USED': 'yes',
                                 'CUSTOM_IMAGE': '{value}',
+                                # This should just be passed through, as JUPYTER_USER is not replaced
+                                'USER': '${JUPYTER_USER}',
+                                # This should render as ${JUPYTER_USER}, as the {{ and }} escape them.
+                                # this matches existing behavior for other replacements elsewhere
+                                'USER_TEST': '${{JUPYTER_USER}}',
                             },
+                            "init_containers": [
+                                {"name": "testing", "image": "{value}"}
+                            ],
                         },
                     },
                 }
@@ -274,7 +282,13 @@ async def test_unlisted_choice_non_string_override():
     await spawner.load_user_options()
 
     assert spawner.image == image
-    assert spawner.environment == {'CUSTOM_IMAGE_USED': 'yes', 'CUSTOM_IMAGE': image}
+    assert spawner.environment == {
+        'CUSTOM_IMAGE_USED': 'yes',
+        'CUSTOM_IMAGE': image,
+        'USER': '${JUPYTER_USER}',
+        'USER_TEST': '${JUPYTER_USER}',
+    }
+    assert spawner.init_containers == [{"name": "testing", "image": image}]
 
 
 async def test_empty_user_options_and_profile_options_api():