refactor if block + pylint + black

Author: christopher-wild

causal_testing/json_front/json_class.py

@@ -203,28 +203,19 @@ def _setup_test(self, causal_test_case: CausalTestCase, test: Mapping) -> tuple[
         minimal_adjustment_set = self.causal_specification.causal_dag.identification(causal_test_case.base_test_case)
         treatment_var = causal_test_case.treatment_variable
         minimal_adjustment_set = minimal_adjustment_set - {treatment_var}
+        estimator_kwargs = {
+            "treatment": treatment_var.name,
+            "treatment_value": causal_test_case.treatment_value,
+            "control_value": causal_test_case.control_value,
+            "adjustment_set": minimal_adjustment_set,
+            "outcome": causal_test_case.outcome_variable.name,
+            "df": causal_test_engine.scenario_execution_data_df,
+            "effect_modifiers": causal_test_case.effect_modifier_configuration,
+        }
         if "formula" in test:
-            estimation_model = test["estimator"](
-                treatment=treatment_var.name,
-                treatment_value=causal_test_case.treatment_value,
-                control_value=causal_test_case.control_value,
-                adjustment_set=minimal_adjustment_set,
-                outcome=causal_test_case.outcome_variable.name,
-                df=causal_test_engine.scenario_execution_data_df,
-                effect_modifiers=causal_test_case.effect_modifier_configuration,
-                formula=test["formula"],
-            )
-        else:
-            estimation_model = test["estimator"](
-                treatment=treatment_var.name,
-                treatment_value=causal_test_case.treatment_value,
-                control_value=causal_test_case.control_value,
-                adjustment_set=minimal_adjustment_set,
-                outcome=causal_test_case.outcome_variable.name,
-                df=causal_test_engine.scenario_execution_data_df,
-                effect_modifiers=causal_test_case.effect_modifier_configuration,
-            )
+            estimator_kwargs["formula"] = test["formula"]
 
+        estimation_model = test["estimator"](**estimator_kwargs)
         return causal_test_engine, estimation_model
 
     def _append_to_file(self, line: str, log_level: int = None):
@@ -235,9 +226,7 @@ def _append_to_file(self, line: str, log_level: int = None):
         is possible to use the inbuilt logging level variables such as logging.INFO and logging.WARNING
         """
         with open(self.output_path, "a", encoding="utf-8") as f:
-            f.write(
-                line
-            )
+            f.write(line)
         if log_level:
             logger.log(level=log_level, msg=line)
 

causal_testing/testing/estimators.py

@@ -320,7 +320,7 @@ def __init__(
             self.formula = formula
         else:
             terms = [treatment] + sorted(list(adjustment_set)) + sorted(list(effect_modifiers))
-            self.formula = f"{outcome} ~ {'+'.join(((terms)))}"
+            self.formula = f"{outcome} ~ {'+'.join(terms)}"
 
         for term in self.effect_modifiers:
             self.adjustment_set.add(term)

tests/json_front_tests/test_json_class.py

@@ -156,8 +156,7 @@ def test_formula_in_json_test(self):
         self.assertIn("test_output ~ test_input", ''.join(temp_out))
 
     def tearDown(self) -> None:
-        pass
-        # remove_temp_dir_if_existent()
+        remove_temp_dir_if_existent()
 
 
 def populate_example(*args, **kwargs):