From 27808dfa43909d679d659555f340dba48cd6574e Mon Sep 17 00:00:00 2001
From: Henrik Andersson <jan@dhigroup.com>
Date: Wed, 22 Nov 2023 15:47:01 +0100
Subject: [PATCH] Set quantity ASAP instead of mutating

---
 modelskill/comparison/_comparison.py | 45 ++++++++++++++++------------
 modelskill/matching.py               |  3 +-
 2 files changed, 27 insertions(+), 21 deletions(-)

diff --git a/modelskill/comparison/_comparison.py b/modelskill/comparison/_comparison.py
index 07b016199..05757326e 100644
--- a/modelskill/comparison/_comparison.py
+++ b/modelskill/comparison/_comparison.py
@@ -221,8 +221,8 @@ def all(self) -> Sequence[str]:
     def parse(
         items: List[str],
         obs_item: str | int | None = None,
-        mod_items: Optional[List[str | int]] = None,
-        aux_items: Optional[List[str | int]] = None,
+        mod_items: Optional[Iterable[str | int]] = None,
+        aux_items: Optional[Iterable[str | int]] = None,
     ) -> ItemSelection:
         """Parse items and return observation, model and auxiliary items
         Default behaviour:
@@ -310,13 +310,14 @@ def _inside_polygon(polygon, xy) -> np.ndarray:
 
 def _matched_data_to_xarray(
     df: pd.DataFrame,
-    obs_item=None,
-    mod_items=None,
-    aux_items=None,
-    name=None,
-    x=None,
-    y=None,
-    z=None,
+    obs_item: int | str | None = None,
+    mod_items: Optional[Iterable[str | int]] = None,
+    aux_items: Optional[Iterable[str | int]] = None,
+    name: Optional[str] = None,
+    x: Optional[float] = None,
+    y: Optional[float] = None,
+    z: Optional[float] = None,
+    quantity: Optional[Quantity] = None,
 ):
     """Convert matched data to accepted xarray.Dataset format"""
     assert isinstance(df, pd.DataFrame)
@@ -346,7 +347,11 @@ def _matched_data_to_xarray(
     else:
         ds.attrs["gtype"] = str(GeometryType.TRACK)
 
-    q = Quantity.undefined()
+    if quantity is None:
+        q = Quantity.undefined()
+    else:
+        q = quantity
+
     ds["Observation"].attrs["long_name"] = q.name
     ds["Observation"].attrs["units"] = q.unit
 
@@ -406,15 +411,16 @@ def __init__(
     @classmethod
     def from_matched_data(
         cls,
-        data,
-        raw_mod_data=None,
-        obs_item=None,
-        mod_items=None,
-        aux_items=None,
-        name=None,
-        x=None,
-        y=None,
-        z=None,
+        data: xr.Dataset | pd.DataFrame,
+        raw_mod_data: Optional[Dict[str, pd.DataFrame]] = None,
+        obs_item: str | int | None = None,
+        mod_items: Optional[Iterable[str | int]] = None,
+        aux_items: Optional[Iterable[str | int]] = None,
+        name: Optional[str] = None,
+        x: Optional[float] = None,
+        y: Optional[float] = None,
+        z: Optional[float] = None,
+        quantity: Optional[Quantity] = None,
     ):
         """Initialize from compared data"""
         if not isinstance(data, xr.Dataset):
@@ -428,6 +434,7 @@ def from_matched_data(
                 x=x,
                 y=y,
                 z=z,
+                quantity=quantity,
             )
         return cls(matched_data=data, raw_mod_data=raw_mod_data)
 
diff --git a/modelskill/matching.py b/modelskill/matching.py
index d7bad17f4..025b380c5 100644
--- a/modelskill/matching.py
+++ b/modelskill/matching.py
@@ -139,9 +139,8 @@ def from_matched(
         x=x,
         y=y,
         z=z,
+        quantity=quantity,
     )
-    if quantity is not None:
-        cmp.quantity = quantity
     return cmp