Some linting

Author: jmafoster1

causal_testing/generation/enum_gen.py

@@ -1,13 +1,40 @@
+"""This module contains the class EnumGen, which allows us to easily create generating uniform distributions from enums."""
+
 from scipy.stats import rv_discrete
+from enum import Enum
+import numpy as np
 
 
 class EnumGen(rv_discrete):
-    def __init__(self, dt: Enum):
-        self.dt = dict(enumerate(dt, 1))
+    """This class allows us to easily create generating uniform distributions from enums. This is helpful for generating concrete test inputs from abstract test cases."""
+
+    def __init__(self, datatype: Enum):
+        self.dt = dict(enumerate(datatype, 1))
         self.inverse_dt = {v: k for k, v in self.dt.items()}
 
-    def ppf(self, q, *args, **kwds):
+    def ppf(self, q):
+        """Percent point function (inverse of `cdf`) at q of the given RV.
+        Parameters
+        ----------
+        q : array_like
+            Lower tail probability.
+        Returns
+        -------
+        k : array_like
+            Quantile corresponding to the lower tail probability, q.
+        """
         return np.vectorize(self.dt.get)(np.ceil(len(self.dt) * q))
 
-    def cdf(self, q, *args, **kwds):
-        return np.vectorize(self.inverse_dt.get)(q) / len(Car)
+    def cdf(self, q):
+        """
+        Cumulative distribution function of the given RV.
+        Parameters
+        ----------
+        q : array_like
+            quantiles
+        Returns
+        -------
+        cdf : ndarray
+            Cumulative distribution function evaluated at `x`
+        """
+        return np.vectorize(self.inverse_dt.get)(q) / len(self.dt)

causal_testing/json_front/json_class.py

@@ -77,11 +77,11 @@ def _create_abstract_test_case(self, test, mutates, effects):
         assert len(test["mutations"]) == 1
         treatment_var = next(self.scenario.variables[v] for v in test["mutations"])
         if not treatment_var.distribution:
-            fitter = Fitter(self.data[var.name], distributions=get_common_distributions())
+            fitter = Fitter(self.data[treatment_var.name], distributions=get_common_distributions())
             fitter.fit()
             (dist, params) = list(fitter.get_best(method="sumsquare_error").items())[0]
-            var.distribution = getattr(scipy.stats, dist)(**params)
-            self._append_to_file(var.name + f" {dist}({params})", logging.INFO)
+            treatment_var.distribution = getattr(scipy.stats, dist)(**params)
+            self._append_to_file(treatment_var.name + f" {dist}({params})", logging.INFO)
 
         abstract_test = AbstractCausalTestCase(
             scenario=self.scenario,