content: #  Copyright (c) ZenML GmbH 2022. All Rights Reserved.
  #
  #  Licensed under the Apache License, Version 2.0 (the "License");
  #  you may not use this file except in compliance with the License.
  #  You may obtain a copy of the License at:
  #
  #       https://www.apache.org/licenses/LICENSE-2.0
  #
  #  Unless required by applicable law or agreed to in writing, software
  #  distributed under the License is distributed on an "AS IS" BASIS,
  #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
  #  or implied. See the License for the specific language governing
  #  permissions and limitations under the License.
  """Base class for all ZenML data validators."""
  
  from typing import Any, ClassVar, Optional, Sequence, Type, cast
  
  from zenml.client import Client
  from zenml.enums import StackComponentType
  from zenml.stack import Flavor, StackComponent
  from zenml.stack.stack_component import StackComponentConfig
  
  
  class BaseDataValidatorConfig(StackComponentConfig):
      """Base config for all data validators."""
  
  
  class BaseDataValidator(StackComponent):
      """Base class for all ZenML data validators."""
  
      NAME: ClassVar[str]
      FLAVOR: ClassVar[Type["BaseDataValidatorFlavor"]]
  
      @property
      def config(self) -> BaseDataValidatorConfig:
          """Returns the config of this data validator.
  
          Returns:
              The config of this data validator.
          """
          return cast(BaseDataValidatorConfig, self._config)
  
      @classmethod
      def get_active_data_validator(cls) -> "BaseDataValidator":
          """Get the data validator registered in the active stack.
  
          Returns:
              The data validator registered in the active stack.
  
          Raises:
              TypeError: if a data validator is not part of the
                  active stack.
          """
          flavor: BaseDataValidatorFlavor = cls.FLAVOR()
          client = Client()
          data_validator = client.active_stack.data_validator
          if not data_validator or not isinstance(data_validator, cls):
              raise TypeError(
                  f"The active stack needs to have a {cls.NAME} data "
                  f"validator component registered to be able to run data validation "
                  f"actions with {cls.NAME}. You can create a new stack with "
                  f"a {cls.NAME} data validator component or update your "
                  f"active stack to add this component, e.g.:\n\n"
                  f"  `zenml data-validator register {flavor.name} "
                  f"--flavor={flavor.name} ...`\n"
                  f"  `zenml stack register <STACK-NAME> -dv {flavor.name} ...`\n"
                  f"  or:\n"
                  f"  `zenml stack update -dv {flavor.name}`\n\n"
              )
  
          return data_validator
  
      def data_profiling(
          self,
          dataset: Any,
          comparison_dataset: Optional[Any] = None,
          profile_list: Optional[Sequence[Any]] = None,
          **kwargs: Any,
      ) -> Any:
          """Analyze one or more datasets and generate a data profile.
  
          This method should be implemented by data validators that support
          analyzing a dataset and generating a data profile (e.g. schema,
          statistical summary, data distribution profile, validation
          rules, data drift reports etc.).
          The method should return a data profile object.
  
          This method also accepts an optional second dataset argument to
          accommodate different categories of data profiling, e.g.:
  
          * profiles generated from a single dataset: schema inference, validation
          rules inference, statistical profiles, data integrity reports
          * differential profiles that need a second dataset for comparison:
          differential statistical profiles, data drift reports
  
          Data validators that support generating multiple categories of data
          profiles should also take in a `profile_list` argument that lists the
          subset of profiles to be generated. If not supplied, the behavior is
          implementation specific, but it is recommended to provide a good default
          (e.g. a single default data profile type may be generated and returned,
          or all available data profiles may be generated and returned as a single
          result).
  
          Args:
              dataset: Target dataset to be profiled.
              comparison_dataset: Optional second dataset to be used for data
                  comparison profiles (e.g data drift reports).
              profile_list: Optional list identifying the categories of data
                  profiles to be generated.
              **kwargs: Implementation specific keyword arguments.
  
          Raises:
              NotImplementedError: if data profiling is not supported by this
                  data validator.
          """
          raise NotImplementedError(
              f"Data profiling is not supported by the {self.__class__} data "
              f"validator."
          )
  
      def data_validation(
          self,
          dataset: Any,
          comparison_dataset: Optional[Any] = None,
          check_list: Optional[Sequence[Any]] = None,
          **kwargs: Any,
      ) -> Any:
          """Run data validation checks on a dataset.
  
          This method should be implemented by data validators that support
          running data quality checks an input dataset (e.g. data integrity
          checks, data drift checks).
  
          This method also accepts an optional second dataset argument to
          accommodate different categories of data validation tests, e.g.:
  
          * single dataset checks: data integrity checks (e.g. missing
          values, conflicting labels, mixed data types etc.)
          * checks that compare two datasets: data drift checks (e.g. new labels,
          feature drift, label drift etc.)
  
          Data validators that support running multiple categories of data
          integrity checks should also take in a `check_list` argument that
          lists the subset of checks to be performed. If not supplied, the
          behavior is implementation specific, but it is recommended to provide a
          good default (e.g. a single default validation check may be performed,
          or all available validation checks may be performed and their results
          returned as a list of objects).
  
          Args:
              dataset: Target dataset to be validated.
              comparison_dataset: Optional second dataset to be used for data
                  comparison checks (e.g data drift checks).
              check_list: Optional list identifying the data checks to
                  be performed.
              **kwargs: Implementation specific keyword arguments.
  
          Raises:
              NotImplementedError: if data validation is not
                  supported by this data validator.
          """
          raise NotImplementedError(
              f"Data validation not implemented for {self}."
          )
  
      def model_validation(
          self,
          dataset: Any,
          model: Any,
          comparison_dataset: Optional[Any] = None,
          check_list: Optional[Sequence[Any]] = None,
          **kwargs: Any,
      ) -> Any:
          """Run model validation checks.
  
          This method should be implemented by data validators that support
          running model validation checks (e.g. confusion matrix validation,
          performance reports, model error analyzes, etc).
  
          Unlike `data_validation`, model validation checks require that a model
          be present as an active component during the validation process.
  
          This method also accepts an optional second dataset argument to
          accommodate different categories of data validation tests, e.g.:
  
          * single dataset tests: confusion matrix validation,
          performance reports, model error analyzes, etc
          * model comparison tests: tests that identify changes in a model
          behavior by comparing how it performs on two different datasets.
  
          Data validators that support running multiple categories of model
          validation checks should also take in a `check_list` argument that
          lists the subset of checks to be performed. If not supplied, the
          behavior is implementation specific, but it is recommended to provide a
          good default (e.g. a single default validation check may be performed,
          or all available validation checks may be performed and their results
          returned as a list of objects).
  
          Args:
              dataset: Target dataset to be validated.
              model: Target model to be validated.
              comparison_dataset: Optional second dataset to be used for model
                  comparison checks (e.g model performance comparison checks).
              check_list: Optional list identifying the model validation checks to
                  be performed.
              **kwargs: Implementation specific keyword arguments.
  
          Raises:
              NotImplementedError: if model validation is not supported by this
                  data validator.
          """
          raise NotImplementedError(
              f"Model validation not implemented for {self}."
          )
  
  
  class BaseDataValidatorFlavor(Flavor):
      """Base class for data validator flavors."""
  
      @property
      def type(self) -> StackComponentType:
          """The type of the component.
  
          Returns:
              The type of the component.
          """
          return StackComponentType.DATA_VALIDATOR
  
      @property
      def config_class(self) -> Type[BaseDataValidatorConfig]:
          """Config class for data validator.
  
          Returns:
              Config class for data validator.
          """
          return BaseDataValidatorConfig
  
      @property
      def implementation_class(self) -> Type[BaseDataValidator]:
          """Implementation for data validator.
  
          Returns:
              Implementation for data validator.
          """
          return BaseDataValidator
  