chore: simplify energy calculator

equinor · Dec 11, 2024 · a210215 · a210215
1 parent 54be8e2
commit a210215
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 18 deletions.
diff --git a/src/libecalc/application/energy_calculator.py b/src/libecalc/application/energy_calculator.py
@@ -11,16 +11,6 @@
 from libecalc.common.variables import ExpressionEvaluator
 from libecalc.core.result import EcalcModelResult
 from libecalc.core.result.emission import EmissionResult
-from libecalc.infrastructure.components import ConsumerSystem
-from libecalc.infrastructure.components import (
-    ElectricityConsumer as ElectricityConsumerDTO,
-)
-from libecalc.infrastructure.components import (
-    FuelConsumer as FuelConsumerDTO,
-)
-from libecalc.infrastructure.components import (
-    GeneratorSet as GeneratorSetDTO,
-)
 
 
 class Context(ComponentEnergyContext):
@@ -78,17 +68,13 @@ def evaluate_energy_usage(self) -> dict[str, EcalcModelResult]:
         energy_components = self._energy_model.get_energy_components()
 
         for energy_component in energy_components:
-            if isinstance(energy_component, ElectricityConsumerDTO | FuelConsumerDTO):
-                self._consumer_results.update(energy_component.evaluate_energy_usage(self._expression_evaluator))
-            elif isinstance(energy_component, GeneratorSetDTO):
+            if hasattr(energy_component, "evaluate_energy_usage"):
                 context = self._get_context(energy_component.id)
                 self._consumer_results.update(
                     energy_component.evaluate_energy_usage(
                         context=context, expression_evaluator=self._expression_evaluator
                     )
                 )
-            elif isinstance(energy_component, ConsumerSystem):
-                self._consumer_results.update(energy_component.evaluate_energy_usage(self._expression_evaluator))
 
         self._consumer_results = Numbers.format_results_to_precision(self._consumer_results, precision=6)
         return self._consumer_results

diff --git a/src/libecalc/infrastructure/components.py b/src/libecalc/infrastructure/components.py
@@ -185,7 +185,9 @@ def get_component_process_type(self) -> ComponentType:
     def get_name(self) -> str:
         return self.name
 
-    def evaluate_energy_usage(self, expression_evaluator: ExpressionEvaluator) -> dict[str, EcalcModelResult]:
+    def evaluate_energy_usage(
+        self, expression_evaluator: ExpressionEvaluator, context: ComponentEnergyContext
+    ) -> dict[str, EcalcModelResult]:
         consumer_results: dict[str, EcalcModelResult] = {}
         consumer = ConsumerEnergyComponent(
             id=self.id,
@@ -248,7 +250,9 @@ def get_component_process_type(self) -> ComponentType:
     def get_name(self) -> str:
         return self.name
 
-    def evaluate_energy_usage(self, expression_evaluator: ExpressionEvaluator) -> dict[str, EcalcModelResult]:
+    def evaluate_energy_usage(
+        self, expression_evaluator: ExpressionEvaluator, context: ComponentEnergyContext
+    ) -> dict[str, EcalcModelResult]:
         consumer_results: dict[str, EcalcModelResult] = {}
         consumer = ConsumerEnergyComponent(
             id=self.id,
@@ -449,7 +453,9 @@ def get_component_process_type(self) -> ComponentType:
     def get_name(self) -> str:
         return self.name
 
-    def evaluate_energy_usage(self, expression_evaluator: ExpressionEvaluator) -> dict[str, EcalcModelResult]:
+    def evaluate_energy_usage(
+        self, expression_evaluator: ExpressionEvaluator, context: ComponentEnergyContext
+    ) -> dict[str, EcalcModelResult]:
         consumer_results = {}
         evaluated_stream_conditions = self.evaluate_stream_conditions(
             expression_evaluator=expression_evaluator,