Add assertions public api methods

dataikuapi/dss/ml.py

@@ -431,6 +431,23 @@ def get_split_params(self):
         """
         return PredictionSplitParamsHandler(self.mltask_settings)
 
+    @property
+    def assertions_params(self):
+        """
+        Retrieves the assertions parameters for this ml task
+
+        :rtype: :class:`dataikuapi.dss.ml.DSSMLAssertionsParams`
+        """
+        return self.get_assertions_params()
+
+    def get_assertions_params(self):
+        """
+        Retrieves the assertions parameters for this ml task
+
+        :rtype: :class:`dataikuapi.dss.ml.DSSMLAssertionsParams`
+        """
+        return DSSMLAssertionsParams(self.mltask_settings["assertionParams"])
+
     def split_ordered_by(self, feature_name, ascending=True):
         """
         Deprecated. Use split_params.set_time_ordering()
@@ -643,6 +660,324 @@ def __repr__(self):
                                                           msg=self._internal_dict["message"])
 
 
+class DSSMLAssertionsParams(object):
+    """
+    Object that represents parameters for all assertions of a ml task
+    Do not create this object directly, use :meth:`DSSPredictionMLTaskSettings.get_assertions_params()` instead
+    """
+
+    def __init__(self, data):
+        self._internal_dict = data
+
+    def get_raw(self):
+        """
+        Gets the raw dictionary of the assertions parameters
+        :rtype: dict
+        """
+        return self._internal_dict
+
+    def get_assertion(self, assertion_name):
+        """
+        Gets a :class:`dataikuapi.dss.ml.DSSMLAssertionParams` representing the parameters of the assertion with the
+        provided name (or None if no assertion has that name)
+        :param str assertion_name: Name of the assertion
+        :rtype: :class:`dataikuapi.dss.ml.DSSMLAssertionParams` or None
+        """
+        for assertion_dict in self._internal_dict["assertions"]:
+            if assertion_dict["name"] == assertion_name:
+                return DSSMLAssertionParams(assertion_dict)
+        return None
+
+    def add_assertion(self, assertion_params):
+        """
+        Adds parameters of an assertion to the assertions parameters of the ml task.
+        Raises a ValueError if an assertion with the same name already exists
+        :param object  assertion_params: A :class:`~dataikuapi.dss.utils.DSSMLAssertionParams` representing parameters of the assertion
+        """
+        if not isinstance(assertion_params, DSSMLAssertionParams):
+            raise ValueError('Wrong type for assertion parameters: {}'.format(type(assertion_params)))
+        self.check_assertion_names_are_uniq(self._internal_dict["assertions"] + [assertion_params._internal_dict])
+
+        self._internal_dict["assertions"].append(assertion_params._internal_dict)
+
+    def delete_assertion(self, assertion_name):
+        """
+        Deletes the assertion parameters of the assertion with the provided name from the `dataikuapi.dss.ml.DSSMLAssertionsParams`
+        Raises a ValueError if no assertion with the provided name was found
+        :param str assertion_name: Name of the assertion
+        """
+        for idx, assertion_dict in enumerate(self._internal_dict["assertions"]):
+            if assertion_dict["name"] == assertion_name:
+                del self._internal_dict["assertions"][idx]
+                return
+        raise ValueError('No assertion found with name: {}'.format(assertion_name))
+
+
+class DSSMLAssertionParams(object):
+    """
+    Object that represents parameters for one assertion
+    Do not create this object directly, use :meth:`dataikuapi.dss.ml.DSSMLAssertionsParams.get_assertion(assertion_name)` or
+    `create_from_parts(name, a_filter, condition)` instead
+    """
+    def __init__(self, data):
+        self._internal_dict = data
+
+    @staticmethod
+    def create_from_parts(name, a_filter, condition):
+        """
+        Creates assertion parameters from name, filter and condition
+
+        :param str name: Name of the assertion
+        :param object a_filter: A :class:`~dataikuapi.dss.utils.DSSFilter`  to select assertion population
+        :param object condition: A :class:`~dataikuapi.dss.ml.DSSMLAssertionCondition` for the assertion to be successful
+
+        :rtype: :class:`dataikuapi.dss.ml.DSSMLAssertionParams`
+        """
+        assertion_params = DSSMLAssertionParams({})
+        assertion_params.name = name
+        assertion_params.filter = a_filter
+        assertion_params.condition = condition
+        return assertion_params
+
+    def get_raw(self):
+        """
+        Gets the raw dictionary of the assertion parameters
+        :rtype: dict
+        """
+        return self._internal_dict
+
+    @property
+    def name(self):
+        """
+        Returns the assertion name
+        :rtype: str
+        """
+        return self._internal_dict["name"]
+
+    @name.setter
+    def name(self, name):
+        self._internal_dict["name"] = name
+
+    @property
+    def filter(self):
+        """
+        Returns the assertion filter
+        :rtype: :class:`dataikuapi.dss.utils.DSSFilter`
+        """
+        return self._internal_dict["filter"]
+
+    @filter.setter
+    def filter(self, filter):
+        self._internal_dict["filter"] = filter
+
+    @property
+    def condition(self):
+        """
+        Returns the assertion condition
+        :rtype: :class:`dataikuapi.dss.ml.DSSMLAssertionCondition`
+        """
+        return DSSMLAssertionCondition(self._internal_dict["assertionCondition"])
+
+    @condition.setter
+    def condition(self, condition):
+        if not isinstance(condition, DSSMLAssertionCondition):
+            raise ValueError('Wrong type for assertion condition: {}.format(type(condition))')
+        self._internal_dict["assertionCondition"] = condition._internal_dict
+
+
+class DSSMLAssertionCondition(object):
+    """
+      Object that represents an assertion condition
+      Do not create this object directly, use :meth:`dataikuapi.dss.ml.DSSMLAssertionParams.condition`, :meth:`dataikuapi.dss.ml.DSSMLAssertionCondition.create_from_expected_class(expected_valid_ratio, expected_class)`
+      or :meth:`dataikuapi.dss.ml.DSSMLAssertionCondition.create_from_expected_range(expected_valid_ratio, expected_range)` instead
+    """
+    def __init__(self, data):
+        self._internal_dict = data
+
+    @staticmethod
+    def create_from_expected_class(expected_valid_ratio, expected_class):
+        """
+        Creates an assertion condition from the expected valid ratio and class
+
+        :param float expected_valid_ratio: Ratio of valid rows needed for the assertion to pass
+        :param str expected_class: Class on which the `expected_valid_ratio` will be calculated
+
+        :rtype: :class:`dataikuapi.dss.ml.DSSMLAssertionCondition`
+        """
+        assertion_condition = DSSMLAssertionCondition({})
+        assertion_condition.expected_valid_ratio = expected_valid_ratio
+        assertion_condition.expected_class = expected_class
+        return assertion_condition
+
+    @staticmethod
+    def create_from_expected_range(expected_valid_ratio, expected_range):
+        """
+        Creates an assertion condition from an expected valid ratio and an expected range
+
+        :param float expected_valid_ratio: Assertion passes if this ratio of rows predicted between expected_min and expected_max is attained
+        :param tuple(float,float) expected_range: Range of values (min, max) where the prediction will be considered as
+        valid for the `expected_valid_ratio`
+
+        :rtype: :class:`dataikuapi.dss.ml.DSSMLAssertionCondition`
+        """
+        assertion_condition = DSSMLAssertionCondition({})
+        assertion_condition.expected_valid_ratio = expected_valid_ratio
+        assertion_condition.expected_range = expected_range
+        return assertion_condition
+
+    def get_raw(self):
+        """
+        Gets the raw dictionary of the condition
+        :rtype: dict
+        """
+        return self._internal_dict
+
+    @property
+    def expected_class(self):
+        """
+        Returns the expected class on which the valid ratio will be calculated
+        :rtype: str
+        """
+        if "expectedClass" in self._internal_dict:
+            return self._internal_dict["expectedClass"]
+        else:
+            return None
+
+    @expected_class.setter
+    def expected_class(self, expected_class):
+        if self.expected_range is not None:
+            raise ValueError("Expected class and expected range can't be both set")
+        self._internal_dict["expectedClass"] = expected_class
+
+    @property
+    def expected_valid_ratio(self):
+        """
+        Returns the ratio of valid rows to exceed for the assertion to pass
+        :rtype: str
+        """
+        return self._internal_dict["successRatio"]
+
+    @expected_valid_ratio.setter
+    def expected_valid_ratio(self, expected_valid_ratio):
+        self._internal_dict["successRatio"] = expected_valid_ratio
+
+    @property
+    def expected_range(self):
+        """
+        Returns the expected range on which the valid ratio will be calculated
+        :rtype: tuple(float,float)
+        """
+        if "expectedMinValue" in self._internal_dict and "expectedMaxValue" in self._internal_dict:
+            return self._internal_dict["expectedMinValue"], self._internal_dict["expectedMaxValue"]
+        else:
+            return None
+
+    @expected_range.setter
+    def expected_range(self, expected_range):
+        if not isinstance(expected_range, tuple):
+            raise ValueError("Expected range needs to be a tuple")
+        if self.expected_class is not None:
+            raise ValueError("Expected class and expected range can't be both set")
+        if expected_range[0] > expected_range:
+            raise ValueError("Expected range needs to be sorted in ascending order. (min value, max value)")
+        self._internal_dict["expectedMinValue"] = expected_range[0]
+        self._internal_dict["expectedMaxValue"] = expected_range[1]
+
+
+class DSSMLAssertionsMetrics(object):
+    """
+    Object that represents the per assertion metrics for all assertions on a trained model
+    Do not create this object directly, use :meth:`dataikuapi.dss.ml.DSSTrainedPredictionModelDetails.get_per_assertion_metrics()` instead
+    """
+    def __init__(self, data):
+        self._internal_dict = data
+
+    def get_raw(self):
+        """
+        Gets the raw dictionary of the assertions metrics
+        :rtype: dict
+        """
+        return self._internal_dict
+
+    def get_per_assertion_metric(self, assertion_name):
+        """
+        Retrieves the metric computed for this trained model for the assertion with the provided name (or None)
+
+        :param str assertion_name: Name of the assertion
+
+        :returns: an object representing assertion metrics
+        :rtype: :class:`dataikuapi.dss.ml.DSSMLAssertionMetric`
+        """
+        for assertion_metric_dict in self._internal_dict["perAssertion"]:
+            if assertion_name == assertion_metric_dict["name"]:
+                return DSSMLAssertionMetric(assertion_metric_dict)
+        return None
+
+    @property
+    def positive_assertion_ratio(self):
+        """
+        Returns the ratio of passing assertions
+        :rtype: float
+        """
+        return self._internal_dict['positiveAssertionsRatio']
+
+
+class DSSMLAssertionMetric(object):
+    """
+    Object that represents the result of an assertion on a trained model
+    Do not create this object directly, use :meth:`dataikuapi.dss.ml.DSSMLPerAssertionMetrics.get_assertion_metric(self, assertion_name)` instead
+    """
+    def __init__(self, data):
+        self._internal_dict = data
+
+    def get_raw(self):
+        """
+        Gets the raw dictionary of metrics of one assertion
+        :rtype: dict
+        """
+        return self._internal_dict
+
+    @property
+    def name(self):
+        """
+        Returns the assertion name
+        :rtype: str
+        """
+        return self._internal_dict["name"]
+
+    @property
+    def result(self):
+        """
+        Returns whether the assertion pass
+        :rtype: bool
+        """
+        return self._internal_dict["result"]
+
+    @property
+    def valid_ratio(self):
+        """
+        Returns the ratio of passing rows in the assertion population
+        :rtype: float
+        """
+        return self._internal_dict["validRatio"]
+
+    @property
+    def nb_matching_rows(self):
+        """
+        Returns the number of rows matching filter
+        :rtype: int
+        """
+        return self._internal_dict["nbMatchingRows"]
+
+    @property
+    def nb_dropped_rows(self):
+        """
+        Returns the number of rows dropped by the model's preprocessing
+        :rtype: int
+        """
+        return self._internal_dict["nbDroppedRows"]
+
+
 class DSSTreeNode(object):
     def __init__(self, tree, i):
         self.tree = tree
@@ -801,6 +1136,22 @@ def get_performance_metrics(self):
                 del clean_snippet[x]
         return clean_snippet
 
+    @property
+    def assertions_metrics(self):
+        """
+        Retrieves assertions metrics computed for this trained model
+        :returns: an object representing assertion metrics
+        :rtype: :class:`dataikuapi.dss.ml.DSSMLAssertionsMetrics`
+        """
+        return self.get_assertions_metrics()
+
+    def get_assertions_metrics(self):
+        """
+        Retrieves assertions metrics computed for this trained model
+        :returns: an object representing assertion metrics
+        :rtype: :class:`dataikuapi.dss.ml.DSSMLAssertionsMetrics`
+        """
+        return DSSMLAssertionsMetrics(self.snippet["assertionsMetrics"])
 
     def get_hyperparameter_search_points(self):
         """