Docstrings and comments

Author: christopher-wild

causal_testing/testing/causal_test_engine.py

@@ -57,8 +57,8 @@ def execute_test_suite(self, test_suite: CausalTestSuite) -> list[CausalTestResu
 
         """Execute a suite of causal tests and return the results in a list
         :param test_suite: CasualTestSuite object
-        :return: test_suite results which contains a list of CausalTestResult objects
-
+        :return: A dictionary where each key is the name of the estimators specified and the values are lists of
+                causal_test_result objects
         """
         if self.scenario_execution_data_df.empty:
             raise Exception("No data has been loaded. Please call load_data prior to executing a causal test case.")

causal_testing/testing/causal_test_suite.py

@@ -3,14 +3,25 @@
 
 class CausalTestSuite(UserDict):
     """
-    A CausalTestSuite is a class holding a dictionary, where the keys are base_test_cases representing the treatment
-    and outcome Variables, and the values are test objects. Test Objects hold a causal_test_case_list which is a
-    list of causal_test_cases which provide control and treatment values, and an iterator of Estimator Class References
+    A CausalTestSuite is an extension of the UserDict class, therefore it behaves as a python dictionary with the added
+    functionality of this class.
+    The dictionary structure should be the keys are base_test_cases representing the treatment and outcome Variables,
+    and the values are test objects. Test Objects hold a causal_test_case_list which is a list of causal_test_cases
+    which provide control and treatment values, and an iterator of Estimator Class References
 
-    This structure can be fed into the CausalTestEngines execute_test_suite function which will iterate over all the
+    This dictionary can be fed into the CausalTestEngines execute_test_suite function which will iterate over all the
     base_test_case's and execute each causal_test_case with each iterator.
     """
 
     def add_test_object(self, base_test_case, causal_test_case_list, estimators, estimate_type: str = "ate"):
+        """
+        A setter object to allow for the easy construction of the dictionary test suite structure
+
+        :param base_test_case: A BaseTestCase object consisting of a treatment variable, outcome variable and effect
+        :param causal_test_case_list: A list of causal test cases to be executed
+        :param estimators: A list of estimators, the execute_test_suite function in the TestEngine will produce a list
+                            of test results for each estimator
+        :param estimate_type: A string which denotes the type of estimate to return
+        """
         test_object = {"tests": causal_test_case_list, "estimators": estimators, "estimate_type": estimate_type}
         self.data[base_test_case] = test_object