From 84e84aee3bef683a408548d69276b6309b24b044 Mon Sep 17 00:00:00 2001
From: TimoDiepers <timo.diepers@gmail.com>
Date: Sun, 22 Sep 2024 19:13:11 +0200
Subject: [PATCH 1/2] allow node obj input in add td to exc util

---
 bw_timex/utils.py | 87 +++++++++++++++++++++++++++++++++++++++++------
 1 file changed, 76 insertions(+), 11 deletions(-)

diff --git a/bw_timex/utils.py b/bw_timex/utils.py
index 2356374..28c61c7 100644
--- a/bw_timex/utils.py
+++ b/bw_timex/utils.py
@@ -357,35 +357,100 @@ def plot_characterized_inventory_as_waterfall(
 def get_exchange(**kwargs) -> Exchange:
     """
     Get an exchange from the database.
+
+    Parameters
+    ----------
+    **kwargs :
+        Arguments to specify an exchange.
+            - input_node: Input node object
+            - input_code: Input node code
+            - input_database: Input node database
+            - output_node: Output node object
+            - output_code: Output node code
+            - output_database: Output node database
+
+    Returns
+    -------
+    Exchange
+        The exchange object matching the criteria.
+
+    Raises
+    ------
+    MultipleResults
+        If multiple exchanges match the criteria.
+    UnknownObject
+        If no exchange matches the criteria.
     """
+
+    # Process input_node if present
+    input_node = kwargs.pop("input_node", None)
+    if input_node:
+        kwargs["input_code"] = input_node.code
+        kwargs["input_database"] = input_node.database
+
+    # Process output_node if present
+    output_node = kwargs.pop("output_node", None)
+    if output_node:
+        kwargs["output_code"] = output_node.code
+        kwargs["output_database"] = output_node.database
+
+    # Map kwargs to database fields
     mapping = {
         "input_code": ExchangeDataset.input_code,
         "input_database": ExchangeDataset.input_database,
         "output_code": ExchangeDataset.output_code,
         "output_database": ExchangeDataset.output_database,
     }
-    qs = ExchangeDataset.select()
+
+    # Build query filters
+    filters = []
     for key, value in kwargs.items():
-        try:
-            qs = qs.where(mapping[key] == value)
-        except KeyError:
-            continue
+        field = mapping.get(key)
+        if field is not None:
+            filters.append(field == value)
 
+    # Execute query with filters
+    qs = ExchangeDataset.select().where(*filters)
     candidates = [Exchange(obj) for obj in qs]
-    if len(candidates) > 1:
+    num_candidates = len(candidates)
+
+    if num_candidates > 1:
         raise MultipleResults(
-            "Found {} results for the given search. Please be more specific or double-check your system model for duplicates.".format(
-                len(candidates)
-            )
+            f"Found {num_candidates} results for the given search. "
+            "Please be more specific or double-check your system model for duplicates."
         )
-    elif not candidates:
-        raise UnknownObject
+    elif num_candidates == 0:
+        raise UnknownObject("No exchange found matching the criteria.")
+
     return candidates[0]
 
 
 def add_temporal_distribution_to_exchange(
     temporal_distribution: TemporalDistribution, **kwargs
 ):
+    """
+    Adds a temporal distribution to an exchange specified by kwargs.
+
+    Parameters
+    ----------
+    temporal_distribution : TemporalDistribution
+        TemporalDistribution to be added to the exchange.
+    **kwargs :
+        Arguments to specify an exchange.
+            - input_node: Input node object
+            - input_id: Input node database ID
+            - input_code: Input node code
+            - input_database: Input node database
+            - output_node: Output node object
+            - output_id: Output node database ID
+            - output_code: Output node code
+            - output_database: Output node database
+
+    Returns
+    -------
+    None
+        The exchange is saved with the temporal distribution.
+    """
     exchange = get_exchange(**kwargs)
     exchange["temporal_distribution"] = temporal_distribution
     exchange.save()

From 0c19c8e1431dd9698aacec5f5afcc56951d0565b Mon Sep 17 00:00:00 2001
From: TimoDiepers <timo.diepers@gmail.com>
Date: Tue, 24 Sep 2024 14:29:14 +0200
Subject: [PATCH 2/2] allow node obj input in add td to exc util

---
 bw_timex/utils.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/bw_timex/utils.py b/bw_timex/utils.py
index 28c61c7..e7fed2b 100644
--- a/bw_timex/utils.py
+++ b/bw_timex/utils.py
@@ -385,14 +385,14 @@ def get_exchange(**kwargs) -> Exchange:
     # Process input_node if present
     input_node = kwargs.pop("input_node", None)
     if input_node:
-        kwargs["input_code"] = input_node.code
-        kwargs["input_database"] = input_node.database
+        kwargs["input_code"] = input_node["code"]
+        kwargs["input_database"] = input_node["database"]
 
     # Process output_node if present
     output_node = kwargs.pop("output_node", None)
     if output_node:
-        kwargs["output_code"] = output_node.code
-        kwargs["output_database"] = output_node.database
+        kwargs["output_code"] = output_node["code"]
+        kwargs["output_database"] = output_node["database"]
 
     # Map kwargs to database fields
     mapping = {