From 88ef4c01c9b1160ac86e18ecb92e386a04580ef5 Mon Sep 17 00:00:00 2001
From: RichieHakim <RichHakim@gmail.com>
Date: Mon, 6 May 2024 19:14:37 -0400
Subject: [PATCH] new recursive_valmap function

---
 bnpm/container_helpers.py | 42 +++++++++++++++++++++++++++++++++++++++
 1 file changed, 42 insertions(+)

diff --git a/bnpm/container_helpers.py b/bnpm/container_helpers.py
index 6abac76..a49df2f 100644
--- a/bnpm/container_helpers.py
+++ b/bnpm/container_helpers.py
@@ -4,6 +4,8 @@
 from typing import Dict, Any, Optional, Union, List, Tuple, Callable, Iterable, Iterator, Type
 import itertools
 
+import toolz
+
 """
 This module is intended to have minimal dependencies.
 It is called by server.py which is intended to run
@@ -475,6 +477,46 @@ def make_serializable_dict(obj, depth=0, max_depth=100, name=None):
     return serializable_dict
 
 
+def recursive_valmap(
+    func: Callable[[Any], Any],
+    d: Dict,
+    depth: int = -1,
+):
+    """
+    Apply a function to all values in a nested dictionary. \n
+    If item is a dictionary, then apply the function to all values in the
+    dictionary. \n
+    Uses toolz.dicttoolz.valmap. \n
+    RH 2024
+
+    Args:
+        func (Callable):
+            Function to apply to all values.
+        d (Dict):
+            Dictionary to apply the function to.
+        factory (Type[Dict]):
+            Type of dictionary to use for the output.
+        depth (int):
+            Maximum depth to apply the function.
+            If -1, then apply to all depths.
+            If 1, then apply to the top level only.
+            If 0, then return the dictionary as is.
+
+    Returns:
+        output (Dict):
+            Dictionary with the function applied to all values.
+    """
+    def helper_recursive_valmap(func, d, depth):
+        if depth == 0:
+            return d
+        return toolz.valmap(
+            lambda x: helper_recursive_valmap(func, x, depth-1) if isinstance(x, dict) else func(x),
+            d
+        )
+
+    return helper_recursive_valmap(func, d, depth)
+
+
 ############################################################
 ################# PARAMETER DICTIONARIES ###################
 ############################################################