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:
  #
  #       http://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.
  """Implementation of the Great Expectations materializers."""
  
  import os
  from typing import TYPE_CHECKING, Any, ClassVar, Dict, Tuple, Type, Union, cast
  
  from great_expectations.checkpoint.types.checkpoint_result import (  # type: ignore[import-untyped]
      CheckpointResult,
  )
  from great_expectations.core import (  # type: ignore[import-untyped]
      ExpectationSuite,
  )
  from great_expectations.core.expectation_validation_result import (  # type: ignore[import-untyped]
      ExpectationSuiteValidationResult,
  )
  from great_expectations.data_context.types.base import (
      CheckpointConfig,
  )
  from great_expectations.data_context.types.resource_identifiers import (
      ExpectationSuiteIdentifier,
      ValidationResultIdentifier,
  )
  from great_expectations.types import (  # type: ignore[import-untyped]
      SerializableDictDot,
  )
  
  from zenml.enums import ArtifactType, VisualizationType
  from zenml.integrations.great_expectations.data_validators.ge_data_validator import (
      GreatExpectationsDataValidator,
  )
  from zenml.materializers.base_materializer import BaseMaterializer
  from zenml.utils import source_utils, yaml_utils
  
  if TYPE_CHECKING:
      from zenml.metadata.metadata_types import MetadataType
  
  ARTIFACT_FILENAME = "artifact.json"
  
  
  class GreatExpectationsMaterializer(BaseMaterializer):
      """Materializer to read/write Great Expectation objects."""
  
      ASSOCIATED_TYPES: ClassVar[Tuple[Type[Any], ...]] = (
          ExpectationSuite,
          CheckpointResult,
      )
      ASSOCIATED_ARTIFACT_TYPE: ClassVar[ArtifactType] = (
          ArtifactType.DATA_ANALYSIS
      )
  
      @staticmethod
      def preprocess_checkpoint_result_dict(
          artifact_dict: Dict[str, Any],
      ) -> None:
          """Pre-processes a GE checkpoint dict before it is used to de-serialize a GE CheckpointResult object.
  
          The GE CheckpointResult object is not fully de-serializable
          due to some missing code in the GE codebase. We need to compensate
          for this by manually converting some of the attributes to
          their correct data types.
  
          Args:
              artifact_dict: A dict containing the GE checkpoint result.
          """
  
          def preprocess_run_result(key: str, value: Any) -> Any:
              if key == "validation_result":
                  return ExpectationSuiteValidationResult(**value)
              return value
  
          artifact_dict["checkpoint_config"] = CheckpointConfig(
              **artifact_dict["checkpoint_config"]
          )
          validation_dict = {}
          for result_ident, results in artifact_dict["run_results"].items():
              validation_ident = (
                  ValidationResultIdentifier.from_fixed_length_tuple(  # type: ignore[no-untyped-call]
                      result_ident.split("::")[1].split("/")
                  )
              )
              validation_results = {
                  result_name: preprocess_run_result(result_name, result)
                  for result_name, result in results.items()
              }
              validation_dict[validation_ident] = validation_results
          artifact_dict["run_results"] = validation_dict
  
      def load(self, data_type: Type[Any]) -> SerializableDictDot:
          """Reads and returns a Great Expectations object.
  
          Args:
              data_type: The type of the data to read.
  
          Returns:
              A loaded Great Expectations object.
          """
          filepath = os.path.join(self.uri, ARTIFACT_FILENAME)
          artifact_dict = yaml_utils.read_json(filepath)
          data_type = source_utils.load(artifact_dict.pop("data_type"))
  
          if data_type is CheckpointResult:
              self.preprocess_checkpoint_result_dict(artifact_dict)
  
          return data_type(**artifact_dict)
  
      def save(self, obj: SerializableDictDot) -> None:
          """Writes a Great Expectations object.
  
          Args:
              obj: A Great Expectations object.
          """
          filepath = os.path.join(self.uri, ARTIFACT_FILENAME)
          artifact_dict = obj.to_json_dict()
          artifact_type = type(obj)
          artifact_dict["data_type"] = (
              f"{artifact_type.__module__}.{artifact_type.__name__}"
          )
          yaml_utils.write_json(filepath, artifact_dict)
  
      def save_visualizations(
          self, data: Union[ExpectationSuite, CheckpointResult]
      ) -> Dict[str, VisualizationType]:
          """Saves visualizations for the given Great Expectations object.
  
          Args:
              data: The Great Expectations object to save visualizations for.
  
          Returns:
              A dictionary of visualization URIs and their types.
          """
          visualizations = {}
  
          if isinstance(data, CheckpointResult):
              result = cast(CheckpointResult, data)
              identifier = next(iter(result.run_results.keys()))
          else:
              suite = cast(ExpectationSuite, data)
              identifier = ExpectationSuiteIdentifier(
                  suite.expectation_suite_name
              )
  
          context = GreatExpectationsDataValidator.get_data_context()
          sites = context.get_docs_sites_urls(identifier)
          for site in sites:
              url = site["site_url"]
              visualizations[url] = VisualizationType.HTML
  
          return visualizations
  
      def extract_metadata(
          self, data: Union[ExpectationSuite, CheckpointResult]
      ) -> Dict[str, "MetadataType"]:
          """Extract metadata from the given Great Expectations object.
  
          Args:
              data: The Great Expectations object to extract metadata from.
  
          Returns:
              The extracted metadata as a dictionary.
          """
          if isinstance(data, CheckpointResult):
              return {
                  "checkpoint_result_name": data.name,
                  "checkpoint_result_passed": data.success,
              }
          elif isinstance(data, ExpectationSuite):
              return {
                  "expectation_suite_name": data.name,
              }
          return {}
  