From 806aadd79eb531172d6bfd0b5f49f81c6852bf55 Mon Sep 17 00:00:00 2001
From: Daniel Nata Nugraha <daniel.nata@flower.ai>
Date: Fri, 31 Jan 2025 08:30:30 +0100
Subject: [PATCH] Move try catch

---
 src/py/flwr/cli/config_utils.py | 17 +++--------------
 src/py/flwr/common/config.py    | 16 ++++++++++++++--
 2 files changed, 17 insertions(+), 16 deletions(-)

diff --git a/src/py/flwr/cli/config_utils.py b/src/py/flwr/cli/config_utils.py
index b19d6b910de3..5975edc1ef8d 100644
--- a/src/py/flwr/cli/config_utils.py
+++ b/src/py/flwr/cli/config_utils.py
@@ -167,20 +167,9 @@ def validate_federation_in_project_config(
 
     # Override the federation configuration if provided
     if overrides:
-        try:
-            overrides_dict = parse_config_args(overrides, flatten=False)
-            federation_config = fuse_dicts(federation_config, overrides_dict)
-        except tomli.TOMLDecodeError as e:
-            typer.secho(
-                f"❌ Error parsing overrides due to invalid format. Ensure that "
-                "your overrides uses supported types of bool, int, string, or "
-                "float, e.g. true/false, \"random string\", or 123 and is formatted "
-                "correctly, e.g. 'key1=value1 key2=value2' or 'key3=value3'.",
-                fg=typer.colors.RED,
-                bold=True,
-            )
-            raise typer.Exit(code=1)
-            
+        overrides_dict = parse_config_args(overrides, flatten=False)
+        federation_config = fuse_dicts(federation_config, overrides_dict)
+
     return federation, federation_config
 
 
diff --git a/src/py/flwr/common/config.py b/src/py/flwr/common/config.py
index 19863f8e3d1b..bbfa985b654b 100644
--- a/src/py/flwr/common/config.py
+++ b/src/py/flwr/common/config.py
@@ -23,6 +23,7 @@
 from typing import IO, Any, Optional, TypeVar, Union, cast, get_args
 
 import tomli
+import typer
 
 from flwr.common.constant import (
     APP_DIR,
@@ -233,8 +234,19 @@ def parse_config_args(
 
             matches = pattern.findall(config_line)
             toml_str = "\n".join(f"{k} = {v}" for k, v in matches)
-            overrides.update(tomli.loads(toml_str))
-            flat_overrides = flatten_dict(overrides) if flatten else overrides
+            try:
+                overrides.update(tomli.loads(toml_str))
+                flat_overrides = flatten_dict(overrides) if flatten else overrides
+            except tomli.TOMLDecodeError as err:
+                typer.secho(
+                    "❌ Error parsing overrides due to invalid format. Ensure that "
+                    "your overrides uses supported types of bool, int, string, or "
+                    'float, e.g. true/false, "random string", or 123 and is formatted '
+                    "correctly, e.g. 'key1=value1 key2=value2' or 'key3=value3'.",
+                    fg=typer.colors.RED,
+                    bold=True,
+                )
+                raise typer.Exit(code=1) from err
 
     return flat_overrides