Upgrade num validation fix

NREL · Feb 14, 2024 · 80047fc · 80047fc
1 parent 6faeeb2
commit 80047fc
Showing 1 changed file with 2 additions and 4 deletions.
diff --git a/buildstock_query/tools/upgrades_analyzer.py b/buildstock_query/tools/upgrades_analyzer.py
@@ -366,10 +366,8 @@ def get_report(self, upgrade_num: Optional[int] = None) -> pd.DataFrame:
         """
 
         self._logic_cache = {}
-        max_upg = len(self.upgrade_names) + 1
-        if upgrade_num is not None:
-            if upgrade_num <= 0 or upgrade_num > max_upg:
-                raise ValueError(f"Invalid upgrade {upgrade_num}. Valid upgrade_num = {list(range(1, max_upg))}.")
+        if upgrade_num not in self.upgrade_names:
+            raise ValueError(f"Invalid upgrade {upgrade_num}. Valid upgrade_num = {self.upgrade_names.keys()}.")
 
         record_dfs = []
         for indx, upgrade_names in self.upgrade_names.items():