From f4791e7ed92faa4fc78d760234000b29509294e9 Mon Sep 17 00:00:00 2001
From: TerminalMan <84923604+SecretiveShell@users.noreply.github.com>
Date: Tue, 24 Sep 2024 02:42:01 +0100
Subject: [PATCH] Cleanup config file loader (#208)

* fix config file loader

* prune nonetype values from config dict

fixes default values not initialising properly

* Utils: Shrink None removal function

It is more concise to use a list and dict collection if necessary
rather than iterating through and checking each value. Tested and
works with Tabby's cases.

Signed-off-by: kingbri <bdashore3@proton.me>

---------

Signed-off-by: kingbri <bdashore3@proton.me>
Co-authored-by: kingbri <bdashore3@proton.me>
---
 common/sampling.py     |  4 ++--
 common/tabby_config.py |  7 +++++--
 common/utils.py        | 13 ++++++++++---
 3 files changed, 17 insertions(+), 7 deletions(-)

diff --git a/common/sampling.py b/common/sampling.py
index e49811d..93d6ea2 100644
--- a/common/sampling.py
+++ b/common/sampling.py
@@ -9,7 +9,7 @@ from loguru import logger
 from pydantic import AliasChoices, BaseModel, Field
 from typing import Dict, List, Optional, Union
 
-from common.utils import unwrap, prune_dict
+from common.utils import filter_none_values, unwrap
 
 
 # Common class for sampler params
@@ -403,7 +403,7 @@ def overrides_from_dict(new_overrides: dict):
     """Wrapper function to update sampler overrides"""
 
     if isinstance(new_overrides, dict):
-        overrides_container.overrides = prune_dict(new_overrides)
+        overrides_container.overrides = filter_none_values(new_overrides)
     else:
         raise TypeError("New sampler overrides must be a dict!")
 
diff --git a/common/tabby_config.py b/common/tabby_config.py
index 219e180..d41cc64 100644
--- a/common/tabby_config.py
+++ b/common/tabby_config.py
@@ -11,7 +11,7 @@ from ruamel.yaml.comments import CommentedMap, CommentedSeq
 from ruamel.yaml.scalarstring import PreservedScalarString
 
 from common.config_models import BaseConfigModel, TabbyConfigModel
-from common.utils import merge_dicts, unwrap
+from common.utils import merge_dicts, filter_none_values, unwrap
 
 yaml = YAML(typ=["rt", "safe"])
 
@@ -33,6 +33,9 @@ class TabbyConfig(TabbyConfigModel):
         if not arguments_dict.get("actions"):
             configs.insert(0, self._from_file(pathlib.Path("config.yml")))
 
+        # Remove None (aka unset) values from the configs and merge them together
+        # This should be less expensive than pruning the entire merged dictionary
+        configs = filter_none_values(configs)
         merged_config = merge_dicts(*configs)
 
         # validate and update config
@@ -135,7 +138,7 @@ class TabbyConfig(TabbyConfigModel):
         """loads config from the provided arguments"""
         config = {}
 
-        config_override = args.get("options", {}).get("config", None)
+        config_override = args.get("config", {}).get("config", None)
         if config_override:
             logger.info("Config file override detected in args.")
             config = self._from_file(pathlib.Path(config_override))
diff --git a/common/utils.py b/common/utils.py
index dfa7e92..97ecaf7 100644
--- a/common/utils.py
+++ b/common/utils.py
@@ -19,10 +19,17 @@ def coalesce(*args):
     return next((arg for arg in args if arg is not None), None)
 
 
-def prune_dict(input_dict: Dict) -> Dict:
-    """Trim out instances of None from a dictionary."""
+def filter_none_values(collection: Union[dict, list]) -> Union[dict, list]:
+    """Remove None values from a collection."""
 
-    return {k: v for k, v in input_dict.items() if v is not None}
+    if isinstance(collection, dict):
+        return {
+            k: filter_none_values(v) for k, v in collection.items() if v is not None
+        }
+    elif isinstance(collection, list):
+        return [filter_none_values(i) for i in collection if i is not None]
+    else:
+        return collection
 
 
 def merge_dict(dict1: Dict, dict2: Dict) -> Dict: