Merge pull request #98 from brightway-lca/improve_add_td_to_exc

Improve util function to add TemporalDistribution to Exchange

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()