4 files changed, 41 insertions, 17 deletions

diff --git a/src/openvic-simulation/dataloader/Dataloader.cpp b/src/openvic-simulation/dataloader/Dataloader.cpp
index c21c6da..4a98bca 100644
--- a/src/openvic-simulation/dataloader/Dataloader.cpp
+++ b/src/openvic-simulation/dataloader/Dataloader.cpp
@@ -1,5 +1,12 @@
 #include "Dataloader.hpp"
 
+#include <openvic-dataloader/csv/Parser.hpp>
+#include <openvic-dataloader/detail/CallbackOStream.hpp>
+#include <openvic-dataloader/v2script/Parser.hpp>
+
+#include "openvic-simulation/GameManager.hpp"
+#include "openvic-simulation/utility/Logger.hpp"
+
 using namespace OpenVic;
 using namespace OpenVic::NodeTools;
 using namespace ovdl;
diff --git a/src/openvic-simulation/dataloader/Dataloader.hpp b/src/openvic-simulation/dataloader/Dataloader.hpp
index e1a2614..f833a2c 100644
--- a/src/openvic-simulation/dataloader/Dataloader.hpp
+++ b/src/openvic-simulation/dataloader/Dataloader.hpp
@@ -1,15 +1,8 @@
 #pragma once
 
 #include <filesystem>
-#include <functional>
-#include <vector>
 
 #include "openvic-simulation/dataloader/NodeTools.hpp"
-#include <openvic-dataloader/csv/Parser.hpp>
-#include <openvic-dataloader/detail/CallbackOStream.hpp>
-#include <openvic-dataloader/v2script/Parser.hpp>
-#include "openvic-simulation/GameManager.hpp"
-#include "openvic-simulation/utility/Logger.hpp"
 
 namespace OpenVic {
 	namespace fs = std::filesystem;
diff --git a/src/openvic-simulation/dataloader/NodeTools.cpp b/src/openvic-simulation/dataloader/NodeTools.cpp
index 4deffe1..5ef09d1 100644
--- a/src/openvic-simulation/dataloader/NodeTools.cpp
+++ b/src/openvic-simulation/dataloader/NodeTools.cpp
@@ -134,6 +134,10 @@ node_callback_t NodeTools::expect_colour(callback_t<colour_t> callback) {
 	};
 }
 
+node_callback_t NodeTools::expect_timespan(callback_t<Timespan> callback) {
+	return expect_int(callback);
+}
+
 node_callback_t NodeTools::expect_date(callback_t<Date> callback) {
 	return expect_identifier(
 		[callback](std::string_view identifier) -> bool {
@@ -234,6 +238,22 @@ node_callback_t NodeTools::expect_length(callback_t<size_t> callback) {
 	};
 }
 
+node_callback_t NodeTools::expect_key(std::string_view key, node_callback_t callback) {
+	return _expect_type<ast::AbstractListNode>(
+		[key, callback](ast::AbstractListNode const& list_node) -> bool {
+			std::vector<ast::NodeUPtr> const& list = list_node._statements;
+			for (ast::NodeUPtr const& sub_node : list_node._statements) {
+				ast::AssignNode const* assign_node = sub_node->cast_to<ast::AssignNode>();
+				if (assign_node != nullptr && assign_node->_name == key) {
+					return callback(&*assign_node->_initializer);
+				}
+			}
+			Logger::error("Failed to find expected key: ", key);
+			return false;
+		}
+	);
+}
+
 node_callback_t NodeTools::expect_dictionary_and_length(length_callback_t length_callback, key_value_callback_t callback) {
 	return expect_list_and_length(length_callback, expect_assign(callback));
 }
@@ -267,14 +287,15 @@ key_value_callback_t NodeTools::dictionary_keys_callback(key_map_t& key_map, boo
 	};
 }
 
-bool NodeTools::check_key_map_counts(key_map_t const& key_map) {
+bool NodeTools::check_key_map_counts(key_map_t& key_map) {
 	bool ret = true;
-	for (key_map_t::value_type const& key_entry : key_map) {
-		dictionary_entry_t const& entry = key_entry.second;
+	for (key_map_t::value_type& key_entry : key_map) {
+		dictionary_entry_t& entry = key_entry.second;
 		if (entry.must_appear() && entry.count < 1) {
 			Logger::error("Mandatory dictionary key not present: ", key_entry.first);
 			ret = false;
 		}
+		entry.count = 0;
 	}
 	return ret;
 }
diff --git a/src/openvic-simulation/dataloader/NodeTools.hpp b/src/openvic-simulation/dataloader/NodeTools.hpp
index 20c8dbd..5ba9d63 100644
--- a/src/openvic-simulation/dataloader/NodeTools.hpp
+++ b/src/openvic-simulation/dataloader/NodeTools.hpp
@@ -30,6 +30,7 @@ namespace OpenVic {
 		node_callback_t expect_uint(callback_t<uint64_t> callback);
 		node_callback_t expect_fixed_point(callback_t<fixed_point_t> callback);
 		node_callback_t expect_colour(callback_t<colour_t> callback);
+		node_callback_t expect_timespan(callback_t<Timespan> callback);
 		node_callback_t expect_date(callback_t<Date> callback);
 		node_callback_t expect_ivec2(callback_t<ivec2_t> callback);
 		node_callback_t expect_fvec2(callback_t<fvec2_t> callback);
@@ -43,6 +44,8 @@ namespace OpenVic {
 		node_callback_t expect_list(node_callback_t callback);
 		node_callback_t expect_length(callback_t<size_t> callback);
 
+		node_callback_t expect_key(std::string_view key, node_callback_t callback);
+
 		node_callback_t expect_dictionary_and_length(length_callback_t length_callback, key_value_callback_t callback);
 		node_callback_t expect_dictionary(key_value_callback_t callback);
 
@@ -74,7 +77,7 @@ namespace OpenVic {
 
 		void add_key_map_entry(key_map_t& key_map, std::string_view key, dictionary_entry_t::expected_count_t expected_count, node_callback_t callback);
 		key_value_callback_t dictionary_keys_callback(key_map_t& key_map, bool allow_other_keys);
-		bool check_key_map_counts(key_map_t const& key_map);
+		bool check_key_map_counts(key_map_t& key_map);
 
 		constexpr struct allow_other_keys_t {} ALLOW_OTHER_KEYS;
 
@@ -169,13 +172,13 @@ namespace OpenVic {
 
 		template<typename T>
 		requires(std::integral<T>)
-		callback_t<uint64_t> assign_variable_callback_uint(std::string_view name, T& var) {
-			return [&var, name](uint64_t val) -> bool {
+		callback_t<uint64_t> assign_variable_callback_uint(T& var) {
+			return [&var](uint64_t val) -> bool {
 				if (val <= static_cast<uint64_t>(std::numeric_limits<T>::max())) {
 					var = val;
 					return true;
 				}
-				Logger::error("Invalid ", name, ": ", val, " (valid range: [0, ",
+				Logger::error("Invalid uint: ", val, " (valid range: [0, ",
 					static_cast<uint64_t>(std::numeric_limits<T>::max()), "])");
 				return false;
 			};
@@ -183,13 +186,13 @@ namespace OpenVic {
 
 		template<typename T>
 		requires(std::signed_integral<T>)
-		callback_t<int64_t> assign_variable_callback_int(std::string_view name, T& var) {
-			return [&var, name](int64_t val) -> bool {
+		callback_t<int64_t> assign_variable_callback_int(T& var) {
+			return [&var](int64_t val) -> bool {
 				if (static_cast<int64_t>(std::numeric_limits<T>::lowest()) <= val && val <= static_cast<int64_t>(std::numeric_limits<T>::max())) {
 					var = val;
 					return true;
 				}
-				Logger::error("Invalid ", name, ": ", val, " (valid range: [",
+				Logger::error("Invalid int: ", val, " (valid range: [",
 					static_cast<int64_t>(std::numeric_limits<T>::lowest()), ", ",
 					static_cast<int64_t>(std::numeric_limits<T>::max()), "])");
 				return false;