From 676f0e065b969222ab257263ae4b1a11f056408a Mon Sep 17 00:00:00 2001
From: Danilo Lessa Bernardineli <danilo.lessa@gmail.com>
Date: Thu, 14 Dec 2023 21:07:18 -0300
Subject: [PATCH] More informative error messages when policies / SUFs have the
 wrong output (#318)

* add error msg

---------

Co-authored-by: Emanuel Lima <emanuellima1@users.noreply.github.com>
---
 cadCAD/configuration/__init__.py |  4 ++--
 cadCAD/engine/simulation.py      | 17 ++++++++++++++---
 testing/test_arg_count.py        |  4 ++--
 3 files changed, 18 insertions(+), 7 deletions(-)

diff --git a/cadCAD/configuration/__init__.py b/cadCAD/configuration/__init__.py
index b7e191f8..3e5d07e4 100644
--- a/cadCAD/configuration/__init__.py
+++ b/cadCAD/configuration/__init__.py
@@ -1,9 +1,9 @@
 from typing import Dict, Callable, List, Tuple
-from pandas.core.frame import DataFrame
+from pandas.core.frame import DataFrame # type: ignore
 from datetime import datetime
 from collections import deque
 from copy import deepcopy
-import pandas as pd
+import pandas as pd # type: ignore
 
 from cadCAD.utils import key_filter
 from cadCAD.configuration.utils import exo_update_per_ts, configs_as_objs
diff --git a/cadCAD/engine/simulation.py b/cadCAD/engine/simulation.py
index 1e7829e5..d97d083e 100644
--- a/cadCAD/engine/simulation.py
+++ b/cadCAD/engine/simulation.py
@@ -55,7 +55,12 @@ def compose(init_reduction_funct, funct_list, val_list):
             return result
 
         col_results = get_col_results(sweep_dict, sub_step, sL, s, funcs)
-        key_set = list(set(list(reduce(lambda a, b: a + b, list(map(lambda x: list(x.keys()), col_results))))))
+        try:
+            reducer_arg = list(map(lambda x: list(x.keys()), col_results))
+        except:
+            raise ValueError("There is a Policy Function that has not properly returned a Dictionary")
+        reducer_function = lambda a, b: a + b
+        key_set = list(set(reduce(reducer_function, reducer_arg)))
         new_dict = {k: [] for k in key_set}
         for d in col_results:
             for k in d.keys():
@@ -146,10 +151,16 @@ def transfer_missing_fields(source, destination):
             for k in source:
                 if k not in destination:
                     destination[k] = source[k]
-            del source
+            del source  
             return destination
 
-        last_in_copy: Dict[str, Any] = transfer_missing_fields(last_in_obj, dict(generate_record(state_funcs)))
+        try:
+            new_state_vars = dict(generate_record(state_funcs))
+        except (ValueError, TypeError):
+            raise ValueError("There is a State Update Function which is not returning an proper tuple")
+
+
+        last_in_copy: Dict[str, Any] = transfer_missing_fields(last_in_obj, new_state_vars)
         last_in_copy: Dict[str, Any] = self.apply_env_proc(sweep_dict, env_processes, last_in_copy)
         last_in_copy['substep'], last_in_copy['timestep'], last_in_copy['run'] = sub_step, time_step, run
 
diff --git a/testing/test_arg_count.py b/testing/test_arg_count.py
index e0ac71e6..556c8cfa 100644
--- a/testing/test_arg_count.py
+++ b/testing/test_arg_count.py
@@ -13,10 +13,10 @@ def test_sufs():
     psubs = [
         {
             'policies': {
-                'p_A': lambda _1, _2, _3, _4: {}
+                'p_A': lambda _1, _2, _3, _4, _5: {}
             },
             'variables': {
-                'v_A': lambda _1, _2, _3, _4, _5: ('v_a', None)
+                'v_A': lambda _1, _2, _3, _4, _5, _6: ('a', 1)
             }
         }
     ]