Source code for evalml.data_checks.utils

"""Utility methods for the data checks in EvalML."""

from evalml.data_checks.data_check_action_code import DataCheckActionCode


[docs]def handle_data_check_action_code(action_code): """Handles data check action codes by either returning the DataCheckActionCode or converting from a str. Args: action_code (str or DataCheckActionCode): Data check action code that needs to be handled. Returns: DataCheckActionCode enum Raises: KeyError: If input is not a valid DataCheckActionCode enum value. ValueError: If input is not a string or DatCheckActionCode object. Examples: >>> assert handle_data_check_action_code("drop_col") == DataCheckActionCode.DROP_COL >>> assert handle_data_check_action_code("DROP_ROWS") == DataCheckActionCode.DROP_ROWS >>> assert handle_data_check_action_code("Impute_col") == DataCheckActionCode.IMPUTE_COL """ if isinstance(action_code, str): try: dcac = DataCheckActionCode._all_values[action_code.upper()] except KeyError: raise KeyError("Action code '{}' does not exist".format(action_code)) return dcac if isinstance(action_code, DataCheckActionCode): return action_code raise ValueError( "`handle_data_check_action_code` was not passed a str or DataCheckActionCode object", )