content: #  Copyright (c) ZenML GmbH 2021. 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.
  
  """Utility functions and classes to run ZenML steps."""
  
  import ast
  import contextlib
  import inspect
  import textwrap
  from collections import Counter
  from typing import (
      TYPE_CHECKING,
      Any,
      Callable,
      Dict,
      Optional,
      Tuple,
      TypeVar,
      Union,
  )
  from uuid import UUID
  
  from pydantic import BaseModel
  from typing_extensions import Annotated
  
  from zenml.artifacts.artifact_config import ArtifactConfig
  from zenml.client import Client
  from zenml.enums import (
      ArtifactSaveType,
      ExecutionStatus,
      MetadataResourceTypes,
  )
  from zenml.exceptions import StepInterfaceError
  from zenml.logger import get_logger
  from zenml.metadata.metadata_types import MetadataType
  from zenml.models import RunMetadataResource
  from zenml.steps.step_context import get_step_context
  from zenml.utils import settings_utils, source_code_utils, typing_utils
  
  if TYPE_CHECKING:
      from zenml.steps import BaseStep
  
  
  logger = get_logger(__name__)
  
  SINGLE_RETURN_OUT_NAME = "output"
  
  
  class OutputSignature(BaseModel):
      """The signature of an output artifact."""
  
      resolved_annotation: Any = None
      artifact_config: Optional[ArtifactConfig] = None
      has_custom_name: bool = False
  
      def get_output_types(self) -> Tuple[Any, ...]:
          """Get all output types that match the type annotation.
  
          Returns:
              All output types that match the type annotation.
          """
          if self.resolved_annotation is Any:
              return ()
  
          if typing_utils.is_union(
              typing_utils.get_origin(self.resolved_annotation)
              or self.resolved_annotation
          ):
              return tuple(
                  type(None)
                  if typing_utils.is_none_type(output_type)
                  else output_type
                  for output_type in get_args(self.resolved_annotation)
              )
          else:
              return (self.resolved_annotation,)
  
  
  def get_args(obj: Any) -> Tuple[Any, ...]:
      """Get arguments of a type annotation.
  
      Example:
          `get_args(Union[int, str]) == (int, str)`
  
      Args:
          obj: The annotation.
  
      Returns:
          The args of the annotation.
      """
      return tuple(
          typing_utils.get_origin(v) or v for v in typing_utils.get_args(obj)
      )
  
  
  def parse_return_type_annotations(
      func: Callable[..., Any],
      enforce_type_annotations: bool = False,
  ) -> Dict[str, OutputSignature]:
      """Parse the return type annotation of a step function.
  
      Args:
          func: The step function.
          enforce_type_annotations: If `True`, raises an exception if a type
              annotation is missing.
  
      Raises:
          RuntimeError: If the output annotation has variable length or contains
              duplicate output names.
          RuntimeError: If type annotations should be enforced and a type
              annotation is missing.
  
      Returns:
          - A dictionary mapping output names to their output signatures.
      """
      signature = inspect.signature(func, follow_wrapped=True)
      return_annotation = signature.return_annotation
      output_name: Optional[str]
  
      # Return type annotated as `None`
      if return_annotation is None:
          return {}
  
      # Return type not annotated -> check whether `None` or `Any` should be used
      if return_annotation is signature.empty:
          if enforce_type_annotations:
              raise RuntimeError(
                  "Missing return type annotation for step function "
                  f"'{func.__name__}'."
              )
          elif has_only_none_returns(func):
              return {}
          else:
              return_annotation = Any
  
      if typing_utils.get_origin(return_annotation) is tuple:
          requires_multiple_artifacts = has_tuple_return(func)
          if requires_multiple_artifacts:
              output_signature: Dict[str, Any] = {}
              args = typing_utils.get_args(return_annotation)
              if args[-1] is Ellipsis:
                  raise RuntimeError(
                      "Variable length output annotations are not allowed."
                  )
              for i, annotation in enumerate(args):
                  resolved_annotation = resolve_type_annotation(annotation)
                  artifact_config = get_artifact_config_from_annotation_metadata(
                      annotation
                  )
                  output_name = artifact_config.name if artifact_config else None
                  has_custom_name = output_name is not None
                  output_name = output_name or f"output_{i}"
                  if output_name in output_signature:
                      raise RuntimeError(f"Duplicate output name {output_name}.")
  
                  output_signature[output_name] = OutputSignature(
                      resolved_annotation=resolved_annotation,
                      artifact_config=artifact_config,
                      has_custom_name=has_custom_name,
                  )
              return output_signature
  
      # Return type is annotated as single value or is a tuple
      resolved_annotation = resolve_type_annotation(return_annotation)
      artifact_config = get_artifact_config_from_annotation_metadata(
          return_annotation
      )
      output_name = artifact_config.name if artifact_config else None
      has_custom_name = output_name is not None
      output_name = output_name or SINGLE_RETURN_OUT_NAME
      return {
          output_name: OutputSignature(
              resolved_annotation=resolved_annotation,
              artifact_config=artifact_config,
              has_custom_name=has_custom_name,
          )
      }
  
  
  def resolve_type_annotation(obj: Any) -> Any:
      """Returns the non-generic class for generic aliases of the typing module.
  
      Example: if the input object is `typing.Dict`, this method will return the
      concrete class `dict`.
  
      Args:
          obj: The object to resolve.
  
      Returns:
          The non-generic class for generic aliases of the typing module.
      """
      origin = typing_utils.get_origin(obj) or obj
  
      if origin is Annotated:
          annotation, *_ = typing_utils.get_args(obj)
          return resolve_type_annotation(annotation)
      elif typing_utils.is_union(origin):
          return obj
  
      return origin
  
  
  def get_artifact_config_from_annotation_metadata(
      annotation: Any,
  ) -> Optional[ArtifactConfig]:
      """Get the artifact config from the annotation metadata of a step output.
  
      Example:
      ```python
      get_output_name_from_annotation_metadata(int)  # None
      get_output_name_from_annotation_metadata(Annotated[int, "name"]  # ArtifactConfig(name="name")
      get_output_name_from_annotation_metadata(Annotated[int, ArtifactConfig(name="name", model_name="foo")]  # ArtifactConfig(name="name", model_name="foo")
      ```
  
      Args:
          annotation: The type annotation.
  
      Raises:
          ValueError: If the annotation is not following the expected format
              or if the name was specified multiple times or is an empty string.
  
      Returns:
          The artifact config.
      """
      if (typing_utils.get_origin(annotation) or annotation) is not Annotated:
          return None
  
      annotation, *metadata = typing_utils.get_args(annotation)
  
      error_message = (
          "Artifact annotation should only contain two elements: the artifact "
          "type, and one of the following: { a static or dynamic name || "
          "an `ArtifactConfig` }, e.g.: "
          "`Annotated[int, 'output_name']` || "
          "`Annotated[int, 'output_{placeholder}']` || "
          "`Annotated[int, ArtifactConfig(name='output_{placeholder}')]` ||"
          "`Annotated[int, ArtifactConfig(name='output_name')]`."
      )
  
      if len(metadata) > 2:
          raise ValueError(error_message)
  
      # Loop over all values to also support legacy annotations of the form
      # `Annotated[int, 'output_name', ArtifactConfig(...)]`
      output_name = None
      artifact_config = None
      for metadata_instance in metadata:
          if isinstance(metadata_instance, str):
              if output_name is not None:
                  raise ValueError(error_message)
              output_name = metadata_instance
          elif isinstance(metadata_instance, ArtifactConfig):
              if artifact_config is not None:
                  raise ValueError(error_message)
              artifact_config = metadata_instance
          else:
              raise ValueError(error_message)
  
      # Consolidate output name
      if artifact_config and artifact_config.name:
          if output_name is not None:
              raise ValueError(error_message)
      elif output_name:
          if not artifact_config:
              artifact_config = ArtifactConfig(name=output_name)
          elif not artifact_config.name:
              artifact_config = artifact_config.model_copy()
              artifact_config.name = output_name
  
      if artifact_config and artifact_config.name == "":
          raise ValueError("Output name cannot be an empty string.")
  
      return artifact_config
  
  
  class ReturnVisitor(ast.NodeVisitor):
      """AST visitor class that can be subclassed to visit function returns."""
  
      def __init__(self, ignore_nested_functions: bool = True) -> None:
          """Initializes a return visitor instance.
  
          Args:
              ignore_nested_functions: If `True`, will skip visiting nested
                  functions.
          """
          self._ignore_nested_functions = ignore_nested_functions
          self._inside_function = False
  
      def _visit_function(
          self, node: Union[ast.FunctionDef, ast.AsyncFunctionDef]
      ) -> None:
          """Visit a (async) function definition node.
  
          Args:
              node: The node to visit.
          """
          if self._ignore_nested_functions and self._inside_function:
              # We're already inside a function definition and should ignore
              # nested functions so we don't want to recurse any further
              return
  
          self._inside_function = True
          self.generic_visit(node)
  
      visit_FunctionDef = _visit_function
      visit_AsyncFunctionDef = _visit_function
  
  
  class OnlyNoneReturnsVisitor(ReturnVisitor):
      """Checks whether a function AST contains only `None` returns."""
  
      def __init__(self) -> None:
          """Initializes a visitor instance."""
          super().__init__()
          self.has_only_none_returns = True
  
      def visit_Return(self, node: ast.Return) -> None:
          """Visit a return statement.
  
          Args:
              node: The return statement to visit.
          """
          if node.value is not None:
              if isinstance(node.value, (ast.Constant, ast.NameConstant)):
                  if node.value.value is None:
                      return
  
              self.has_only_none_returns = False
  
  
  class TupleReturnVisitor(ReturnVisitor):
      """Checks whether a function AST contains tuple returns."""
  
      def __init__(self) -> None:
          """Initializes a visitor instance."""
          super().__init__()
          self.has_tuple_return = False
  
      def visit_Return(self, node: ast.Return) -> None:
          """Visit a return statement.
  
          Args:
              node: The return statement to visit.
          """
          if isinstance(node.value, ast.Tuple) and len(node.value.elts) > 1:
              self.has_tuple_return = True
  
  
  def has_tuple_return(func: Callable[..., Any]) -> bool:
      """Checks whether a function returns multiple values.
  
      Multiple values means that the `return` statement is followed by a tuple
      (with or without brackets).
  
      Example:
      ```python
      def f1():
        return 1, 2
  
      def f2():
        return (1, 2)
  
      def f3():
        var = (1, 2)
        return var
  
      has_tuple_return(f1)  # True
      has_tuple_return(f2)  # True
      has_tuple_return(f3)  # False
      ```
  
      Args:
          func: The function to check.
  
      Returns:
          Whether the function returns multiple values.
      """
      source = textwrap.dedent(source_code_utils.get_source_code(func))
      tree = ast.parse(source)
  
      visitor = TupleReturnVisitor()
      visitor.visit(tree)
  
      return visitor.has_tuple_return
  
  
  def has_only_none_returns(func: Callable[..., Any]) -> bool:
      """Checks whether a function contains only `None` returns.
  
      A `None` return could be either an explicit `return None` or an empty
      `return` statement.
  
      Example:
      ```python
      def f1():
        return None
  
      def f2():
        return
  
      def f3(condition):
        if condition:
          return None
        else:
          return 1
  
      has_only_none_returns(f1)  # True
      has_only_none_returns(f2)  # True
      has_only_none_returns(f3)  # False
      ```
  
      Args:
          func: The function to check.
  
      Returns:
          Whether the function contains only `None` returns.
      """
      source = textwrap.dedent(source_code_utils.get_source_code(func))
      tree = ast.parse(source)
  
      visitor = OnlyNoneReturnsVisitor()
      visitor.visit(tree)
  
      return visitor.has_only_none_returns
  
  
  def log_step_metadata(
      metadata: Dict[str, "MetadataType"],
      step_name: Optional[str] = None,
      pipeline_name_id_or_prefix: Optional[Union[str, UUID]] = None,
      run_id: Optional[str] = None,
  ) -> None:
      """Logs step metadata.
  
      Args:
          metadata: The metadata to log.
          step_name: The name of the step to log metadata for. Can be omitted
              when being called inside a step.
          pipeline_name_id_or_prefix: The name of the pipeline to log metadata
              for. Can be omitted when being called inside a step.
          run_id: The ID of the run to log metadata for. Can be omitted when
              being called inside a step.
  
      Raises:
          ValueError: If no step name is provided and the function is not called
              from within a step or if no pipeline name or ID is provided and
              the function is not called from within a step.
      """
      logger.warning(
          "The `log_step_metadata` function is deprecated and will soon be "
          "removed. Please use `log_metadata` instead."
      )
  
      step_context = None
      if not step_name:
          with contextlib.suppress(RuntimeError):
              step_context = get_step_context()
              step_name = step_context.step_name
      # not running within a step and no user-provided step name
      if not step_name:
          raise ValueError(
              "No step name provided and you are not running "
              "within a step. Please provide a step name."
          )
  
      client = Client()
      if step_context:
          step_run_id = step_context.step_run.id
      elif run_id:
          step_run_id = UUID(int=int(run_id))
      else:
          if not pipeline_name_id_or_prefix:
              raise ValueError(
                  "No pipeline name or ID provided and you are not running "
                  "within a step. Please provide a pipeline name or ID, or "
                  "provide a run ID."
              )
          pipeline_run = client.get_pipeline(
              name_id_or_prefix=pipeline_name_id_or_prefix,
          ).last_run
          step_run_id = pipeline_run.steps[step_name].id
      client.create_run_metadata(
          metadata=metadata,
          resources=[
              RunMetadataResource(
                  id=step_run_id, type=MetadataResourceTypes.STEP_RUN
              )
          ],
      )
  
  
  def run_as_single_step_pipeline(
      __step: "BaseStep", *args: Any, **kwargs: Any
  ) -> Any:
      """Runs the step as a single step pipeline.
  
      - All inputs that are not JSON serializable will be uploaded to the
      artifact store before the pipeline is being executed.
      - All output artifacts of the step will be loaded using the materializer
      that was used to store them.
  
      Args:
          *args: Entrypoint function arguments.
          **kwargs: Entrypoint function keyword arguments.
  
      Raises:
          RuntimeError: If the step execution failed.
          StepInterfaceError: If the arguments to the entrypoint function are
              invalid.
  
      Returns:
          The output of the step entrypoint function.
      """
      from zenml import ExternalArtifact, pipeline
      from zenml.config.base_settings import BaseSettings
      from zenml.pipelines.run_utils import (
          wait_for_pipeline_run_to_finish,
      )
  
      logger.info(
          "Running single step pipeline to execute step `%s`", __step.name
      )
  
      try:
          validated_arguments = (
              inspect.signature(__step.entrypoint)
              .bind(*args, **kwargs)
              .arguments
          )
      except TypeError as e:
          raise StepInterfaceError(
              "Invalid step function entrypoint arguments. Check out the "
              "error above for more details."
          ) from e
  
      inputs: Dict[str, Any] = {}
      for key, value in validated_arguments.items():
          try:
              __step.entrypoint_definition.validate_input(key=key, value=value)
              inputs[key] = value
          except Exception:
              inputs[key] = ExternalArtifact(value=value)
  
      orchestrator = Client().active_stack.orchestrator
  
      pipeline_settings: Any = {}
      if "synchronous" in type(orchestrator.config).model_fields:
          # Make sure the orchestrator runs sync so we stream the logs
          key = settings_utils.get_stack_component_setting_key(orchestrator)
          pipeline_settings[key] = BaseSettings(synchronous=True)
  
      @pipeline(name=__step.name, enable_cache=False, settings=pipeline_settings)
      def single_step_pipeline() -> None:
          __step(**inputs)
  
      run = single_step_pipeline.with_options(unlisted=True)()
      assert run
      run = wait_for_pipeline_run_to_finish(run.id)
  
      if run.status != ExecutionStatus.COMPLETED:
          raise RuntimeError("Failed to execute step %s.", __step.name)
  
      # 4. Load output artifacts
      step_run = next(iter(run.steps.values()))
      outputs = [
          artifact_version.load()
          for output_name in step_run.config.outputs.keys()
          for artifact_version in step_run.outputs[output_name]
          if artifact_version.save_type == ArtifactSaveType.STEP_OUTPUT
      ]
  
      if len(outputs) == 0:
          return None
      elif len(outputs) == 1:
          return outputs[0]
      else:
          return tuple(outputs)
  
  
  T = TypeVar("T")
  
  
  def get_unique_step_output_names(
      step_outputs: Dict[Tuple[str, str], T],
  ) -> Dict[Tuple[str, str], Tuple[T, str]]:
      """Get unique step output names.
  
      Given a dictionary of step outputs indexed by (invocation_id, output_name),
      where the value is an arbitrary context object, return the same dictionary
      complemented with unique step output names.
  
      If an output name is repeated, the context object is prepended to the
      output_name to make it unique.
  
      Args:
          step_outputs: The step outputs as arbitrary objects indexed by
              (invocation_id, output_name).
  
      Returns:
          The input dictionary complemented with unique step output names.
      """
      output_name_count = Counter(
          output_name for _, output_name in step_outputs.keys()
      )
  
      unique_step_output_mapping = {}
      for invocation_id, output_name in step_outputs.keys():
          if output_name_count[output_name] > 1:
              unique_step_output_name = f"{invocation_id}.{output_name}"
          else:
              unique_step_output_name = output_name
          unique_step_output_mapping[invocation_id, output_name] = (
              step_outputs[invocation_id, output_name],
              unique_step_output_name,
          )
      return unique_step_output_mapping
  