Codacy cleanup

Author: christopher-wild

causal_testing/json_front/json_class.py

@@ -1,10 +1,9 @@
 from pathlib import Path
 
 from abc import ABC
+import json
 from fitter import Fitter, get_common_distributions
 import pandas as pd
-
-import json
 import scipy
 
 from causal_testing.specification.variable import Input, Output, Meta

examples/poisson/causal_test_setup.py

@@ -9,6 +9,7 @@
 from causal_testing.testing.causal_test_outcome import ExactValue, Positive, Negative, NoEffect, CausalTestOutcome, \
     CausalTestResult
 from causal_testing.json_front.json_class import JsonUtility
+from causal_testing.testing.estimators import Estimator
 
 
 class WidthHeightEstimator(LinearRegressionEstimator):
@@ -153,13 +154,13 @@ def get_args() -> argparse.Namespace:
 class MyJsonUtility(JsonUtility):
     """Extension of JsonUtility class to add modelling assumptions to the estimator instance"""
 
-    def add_modelling_assumptions(self, estimator: LinearRegressionEstimator):
+    def add_modelling_assumptions(self, estimation_model: Estimator):
         # Add squared intensity term as a modelling assumption if intensity is the treatment of the test
-        if "intensity" in estimator.treatment[0]:
-            estimator.add_squared_term_to_df("intensity")
-        if isinstance(estimator, WidthHeightEstimator):
-            estimator.add_product_term_to_df("width", "intensity")
-            estimator.add_product_term_to_df("height", "intensity")
+        if "intensity" in estimation_model.treatment[0]:
+            estimation_model.add_squared_term_to_df("intensity")
+        if isinstance(estimation_model, WidthHeightEstimator):
+            estimation_model.add_product_term_to_df("width", "intensity")
+            estimation_model.add_product_term_to_df("height", "intensity")
 
 
 if __name__ == "__main__":