Group local instance attributes to dataclasses

Author: christopher-wild

causal_testing/json_front/json_class.py

@@ -1,10 +1,12 @@
 """This module contains the JsonUtility class, details of using this class can be found here:
 https://causal-testing-framework.readthedocs.io/en/latest/json_front_end.html"""
+
 import argparse
 import json
 import logging
 
 from abc import ABC
+from dataclasses import dataclass
 from pathlib import Path
 
 import pandas as pd
@@ -42,49 +44,45 @@ class JsonUtility(ABC):
     """
 
     def __init__(self, log_path):
-        self.json_path = None
-        self.dag_path = None
-        self.data_path = None
-        self.inputs = None
-        self.outputs = None
-        self.metas = None
+        self.paths = None
+        self.variables = None
         self.data = None
         self.test_plan = None
         self.modelling_scenario = None
         self.causal_specification = None
         self.setup_logger(log_path)
 
-    def set_path(self, json_path: str, dag_path: str, data_path: str):
+    def set_paths(self, json_path: str, dag_path: str, data_path: str):
         """
         Takes a path of the directory containing all scenario specific files and creates individual paths for each file
         :param json_path: string path representation to .json file containing test specifications
         :param dag_path: string path representation to the .dot file containing the Causal DAG
         :param data_path: string path representation to the data file
-        :returns:
-            - json_path -
-            - dag_path -
-            - data_path -
         """
-        self.json_path = Path(json_path)
-        self.dag_path = Path(dag_path)
-        self.data_path = Path(data_path)
+        self.paths = JsonClassPaths(
+            json_path=json_path,
+            dag_path=dag_path,
+            data_path=data_path
+        )
 
-    def set_variables(self, inputs: dict, outputs: dict, metas: dict):
+    def set_variables(self, inputs: list[dict], outputs: list[dict], metas: list[dict]):
         """Populate the Causal Variables
         :param inputs:
         :param outputs:
         :param metas:
         """
-        self.inputs = [Input(i["name"], i["type"], i["distribution"]) for i in inputs]
-        self.outputs = [Output(i["name"], i["type"], i.get("distribution", None)) for i in outputs]
-        self.metas = [Meta(i["name"], i["type"], i["populate"]) for i in metas] if metas else []
 
+        self.variables = CausalVariables(
+            inputs=inputs,
+            outputs=outputs,
+            metas=metas
+                                          )
     def setup(self):
         """Function to populate all the necessary parts of the json_class needed to execute tests"""
-        self.modelling_scenario = Scenario(self.inputs + self.outputs + self.metas, None)
+        self.modelling_scenario = Scenario(self.variables.inputs + self.variables.outputs + self.variables.metas, None)
         self.modelling_scenario.setup_treatment_variables()
         self.causal_specification = CausalSpecification(
-            scenario=self.modelling_scenario, causal_dag=CausalDAG(self.dag_path)
+            scenario=self.modelling_scenario, causal_dag=CausalDAG(self.paths.dag_path)
         )
         self._json_parse()
         self._populate_metas()
@@ -139,20 +137,20 @@ def _execute_tests(self, concrete_tests, estimators, test, f_flag):
 
     def _json_parse(self):
         """Parse a JSON input file into inputs, outputs, metas and a test plan"""
-        with open(self.json_path, encoding="utf-8") as f:
+        with open(self.paths.json_path, encoding="utf-8") as f:
             self.test_plan = json.load(f)
 
-        self.data = pd.read_csv(self.data_path)
+        self.data = pd.read_csv(self.paths.data_path)
 
     def _populate_metas(self):
         """
         Populate data with meta-variable values and add distributions to Causal Testing Framework Variables
         """
 
-        for meta in self.metas:
+        for meta in self.variables.metas:
             meta.populate(self.data)
 
-        for var in self.metas + self.outputs:
+        for var in self.variables.metas + self.variables.outputs:
             if not var.distribution:
                 fitter = Fitter(self.data[var.name], distributions=get_common_distributions())
                 fitter.fit()
@@ -202,7 +200,7 @@ def _setup_test(self, causal_test_case: CausalTestCase, estimator: Estimator) ->
                 - causal_test_engine - Test Engine instance for the test being run
                 - estimation_model - Estimator instance for the test being run
         """
-        data_collector = ObservationalDataCollector(self.modelling_scenario, self.data_path)
+        data_collector = ObservationalDataCollector(self.modelling_scenario, self.paths.data_path)
         causal_test_engine = CausalTestEngine(self.causal_specification, data_collector, index_col=0)
         minimal_adjustment_set = self.causal_specification.causal_dag.identification(causal_test_case.base_test_case)
         treatment_var = causal_test_case.treatment_variable
@@ -273,3 +271,48 @@ def get_args(test_args=None) -> argparse.Namespace:
             required=True,
         )
         return parser.parse_args(test_args)
+
+
+@dataclass
+class JsonClassPaths:
+    """
+    A dataclass that converts strings of paths to Path objects for use in the JsonUtility class
+    :param json_path: string path representation to .json file containing test specifications
+    :param dag_path: string path representation to the .dot file containing the Causal DAG
+    :param data_path: string path representation to the data file
+    """
+
+    json_path: Path
+    dag_path: Path
+    data_path: Path
+
+    def __init__(
+            self,
+            json_path: str,
+            dag_path: str,
+            data_path: str
+    ):
+        self.json_path = Path(json_path)
+        self.dag_path = Path(dag_path)
+        self.data_path = Path(data_path)
+
+
+@dataclass()
+class CausalVariables:
+    """
+    A dataclass that converts
+    """
+
+    inputs: list[Input]
+    outputs: list[Output]
+    metas: list[Meta]
+
+    def __init__(
+            self,
+            inputs: list[dict],
+            outputs: list[dict],
+            metas: list[dict]
+    ):
+        self.inputs = [Input(i["name"], i["type"], i["distribution"]) for i in inputs]
+        self.outputs = [Output(i["name"], i["type"], i.get("distribution", None)) for i in outputs]
+        self.metas = [Meta(i["name"], i["type"], i["populate"]) for i in metas] if metas else []

tests/json_front_tests/test_json_class.py

@@ -35,29 +35,29 @@ def setUp(self) -> None:
         self.output_dict_list = [{"name": "test_output", "type": float}]
         self.meta_dict_list = [{"name": "test_meta", "type": float, "populate": populate_example}]
         self.json_class.set_variables(self.input_dict_list, self.output_dict_list, None)
-        self.json_class.set_path(self.json_path, self.dag_path, self.data_path)
+        self.json_class.set_paths(self.json_path, self.dag_path, self.data_path)
 
     def test_setting_paths(self):
-        self.assertEqual(self.json_class.json_path, Path(self.json_path))
-        self.assertEqual(self.json_class.dag_path, Path(self.dag_path))
-        self.assertEqual(self.json_class.data_path, Path(self.data_path))
+        self.assertEqual(self.json_class.paths.json_path, Path(self.json_path))
+        self.assertEqual(self.json_class.paths.dag_path, Path(self.dag_path))
+        self.assertEqual(self.json_class.paths.data_path, Path(self.data_path))
 
     def test_set_inputs(self):
         ctf_input = [Input("test_input", float, self.example_distribution)]
-        self.assertEqual(ctf_input[0].name, self.json_class.inputs[0].name)
-        self.assertEqual(ctf_input[0].datatype, self.json_class.inputs[0].datatype)
-        self.assertEqual(ctf_input[0].distribution, self.json_class.inputs[0].distribution)
+        self.assertEqual(ctf_input[0].name, self.json_class.variables.inputs[0].name)
+        self.assertEqual(ctf_input[0].datatype, self.json_class.variables.inputs[0].datatype)
+        self.assertEqual(ctf_input[0].distribution, self.json_class.variables.inputs[0].distribution)
 
     def test_set_outputs(self):
         ctf_output = [Output("test_output", float)]
-        self.assertEqual(ctf_output[0].name, self.json_class.outputs[0].name)
-        self.assertEqual(ctf_output[0].datatype, self.json_class.outputs[0].datatype)
+        self.assertEqual(ctf_output[0].name, self.json_class.variables.outputs[0].name)
+        self.assertEqual(ctf_output[0].datatype, self.json_class.variables.outputs[0].datatype)
 
     def test_set_metas(self):
         self.json_class.set_variables(self.input_dict_list, self.output_dict_list, self.meta_dict_list)
         ctf_meta = [Meta("test_meta", float, populate_example)]
-        self.assertEqual(ctf_meta[0].name, self.json_class.metas[0].name)
-        self.assertEqual(ctf_meta[0].datatype, self.json_class.metas[0].datatype)
+        self.assertEqual(ctf_meta[0].name, self.json_class.variables.metas[0].name)
+        self.assertEqual(ctf_meta[0].datatype, self.json_class.variables.metas[0].datatype)
 
     def test_argparse(self):
         args = self.json_class.get_args(["--data_path=data.csv", "--dag_path=dag.dot", "--json_path=tests.json"])