Another approach to dynamically decorating a class

rules-engine/tests/test_rules_engine/test_examples.py

@@ -29,15 +29,15 @@
 INPUT_DATA = filter(lambda d: d == "breslow", next(os.walk(ROOT_DIR))[1])
 # INPUT_DATA = filter(lambda d: d == "cali", next(os.walk(ROOT_DIR))[1])
 
-
+# Approach number 1:
 # Extend NG Billing Record Input to capture whole home heat loss input from example data
-class NaturalGasBillingRecordExampleInput(NaturalGasBillingRecordInput):
-    whole_home_heat_loss_rate: float
-
-
-# Overload input to capture whole home heat loss input from example data
-class NaturalGasBillingExampleInput(NaturalGasBillingInput):
-    records: List[NaturalGasBillingRecordExampleInput]
+# class NaturalGasBillingRecordExampleInput(NaturalGasBillingRecordInput):
+#    whole_home_heat_loss_rate: float
+#
+# Then overload NG Billing Input to contain new NG Billing Record Example Input subclass
+# However mypy doesn't like this overloading of the "records" attribute
+# class NaturalGasBillingExampleInput(NaturalGasBillingInput):
+#     records: List[NaturalGasBillingRecordExampleInput]
 
 
 class Summary(SummaryInput, SummaryOutput):
@@ -46,17 +46,37 @@ class Summary(SummaryInput, SummaryOutput):
 
 class Example(BaseModel):
     summary: Summary
-    natural_gas_usage: NaturalGasBillingExampleInput
+    natural_gas_usage: NaturalGasBillingInput
     temperature_data: TemperatureInput
 
 
+# Approach number 2:
+# Simply decorate the existing Natural Gas Billing Record Input class with a dynamic attribute for UA
+# This approach leaves the original class interface as-is, but adds this one attribute for testing purposes
+# Had to use the class __dict__ method of access to avoid recursion on get and set
+def whole_home_heat_loss_rate_get(self):
+    return self.__dict__["whole_home_heat_loss_rate"]
+
+
+def whole_home_heat_loss_rate_set(self, value):
+    self.__dict__["whole_home_heat_loss_rate"] = value
+
+
+# Approach number 2:
+# This code runs fine, by mypy raises errors as NaturalGasBillingRecordInput has no heat loss property
+# Extend the existing class to add a whole home heat loss rate property, with get and set methods
+NaturalGasBillingRecordInput.whole_home_heat_loss_rate = property(
+    whole_home_heat_loss_rate_get, whole_home_heat_loss_rate_set
+)
+
+
 def load_summary(folder: str) -> Summary:
     with open(ROOT_DIR / folder / "summary.json") as f:
         d = json.load(f)
         return Summary(**d)
 
 
-def load_natural_gas(folder: str) -> NaturalGasBillingExampleInput:
+def load_natural_gas(folder: str) -> NaturalGasBillingInput:
     records = []
 
     with open(ROOT_DIR / folder / "natural-gas.csv") as f:
@@ -70,23 +90,27 @@ def load_natural_gas(folder: str) -> NaturalGasBillingExampleInput:
                 inclusion_override = int(inclusion_override)
             ua = row["ua"]
             if ua == "":
-                whole_home_heat_loss_rate = float(0)
+                whole_home_heat_loss_rate = None
             else:
                 whole_home_heat_loss_rate = float(ua)
 
-            item = NaturalGasBillingRecordExampleInput(
+            item = NaturalGasBillingRecordInput(
                 period_start_date=datetime.strptime(
                     row["start_date"], "%m/%d/%Y"
                 ).date(),
                 period_end_date=datetime.strptime(row["end_date"], "%m/%d/%Y").date(),
                 usage_therms=row["usage"],
                 inclusion_override=inclusion_override,
-                whole_home_heat_loss_rate=whole_home_heat_loss_rate,
+                # Approach number 1:
+                # whole_home_heat_loss_rate=_whole_home_heat_loss_rate,
             )
+            # Approach number 2:
+            # Extend NaturalGasBillingRecordsInput to track whole home heat loss example input
+            item.whole_home_heat_loss_rate = whole_home_heat_loss_rate
 
             records.append(item)
 
-    return NaturalGasBillingExampleInput(records=records)
+    return NaturalGasBillingInput(records=records)
 
 
 def load_temperature_data(weather_station: str) -> TemperatureInput: