content: #  Copyright (c) ZenML GmbH 2024. 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 for python functions."""
  
  import inspect
  import os
  from contextlib import contextmanager
  from pathlib import Path
  from typing import Any, Callable, Iterator, List, Tuple, TypeVar, Union
  
  import click
  
  from zenml.logger import get_logger
  from zenml.utils.string_utils import random_str
  
  F = TypeVar("F", bound=Callable[..., None])
  
  logger = get_logger(__name__)
  
  _CLI_WRAPPED_SCRIPT_TEMPLATE_HEADER = """
  from zenml.utils.function_utils import _cli_wrapped_function
  
  import sys
  sys.path.append(r"{func_path}")
  
  from {func_module} import {func_name} as step_function
  
  if unwrapped_entrypoint:=getattr(step_function, "unwrapped_entrypoint", None):
      func = _cli_wrapped_function(unwrapped_entrypoint)
  else:
      func = _cli_wrapped_function(step_function.entrypoint)
  """
  _CLI_WRAPPED_MAINS = {
      "accelerate": """
  if __name__=="__main__":
      from accelerate import Accelerator
      import cloudpickle as pickle
      accelerator = Accelerator()
      ret = func(standalone_mode=False)
      if accelerator.is_main_process:
          pickle.dump(ret, open(r"{output_file}", "wb"))
  """
  }
  _ALLOWED_TYPES = (str, int, float, bool, Path)
  _ALLOWED_COLLECTIONS = (tuple,)
  _CLICK_TYPES_MAPPER = {
      str: click.STRING,
      int: click.INT,
      float: click.FLOAT,
      bool: click.BOOL,
      Path: click.STRING,
      None: click.STRING,
  }
  
  
  def _cli_arg_name(arg_name: str) -> str:
      return arg_name.replace("_", "-")
  
  
  def _is_valid_collection_arg(arg_type: Any) -> bool:
      """Check if the given argument type is a valid collection type.
  
      Args:
          arg_type: The type to check.
  
      Returns:
          True if the argument type is a valid collection type, False otherwise.
      """
      if getattr(arg_type, "__origin__", None) in _ALLOWED_COLLECTIONS:
          if arg_type.__args__[0] not in _ALLOWED_TYPES:
              return False
          return True
      return False
  
  
  def _is_valid_optional_arg(arg_type: Any) -> bool:
      """Check if the given argument type is a valid Optional type.
  
      A valid Optional type is defined as a Union with two arguments, where:
      - The first argument is either an allowed type or a valid collection type.
      - The second argument is the NoneType.
  
      Args:
          arg_type: The type to check.
  
      Returns:
          True if the argument type is a valid Optional type, False otherwise.
      """
      if (
          getattr(arg_type, "_name", None) == "Optional"
          and getattr(arg_type, "__origin__", None) == Union
      ):
          if args := getattr(arg_type, "__args__", None):
              if len(args) != 2:
                  return False
              if (
                  args[0] not in _ALLOWED_TYPES
                  and not _is_valid_collection_arg(args[0])
              ) or args[1] is not type(None):
                  return False
          return True
      return False
  
  
  def _cli_wrapped_function(func: F) -> F:
      """Create a decorator to generate the CLI-wrapped function.
  
      Args:
          func: The function to decorate.
  
      Returns:
          The inner decorator.
  
      Raises:
          ValueError: If the function arguments are not valid.
      """
      options: List[Any] = []
      fullargspec = inspect.getfullargspec(func)
      if fullargspec.defaults is not None:
          defaults = [None] * (
              len(fullargspec.args) - len(fullargspec.defaults)
          ) + list(fullargspec.defaults)
      else:
          defaults = [None] * len(fullargspec.args)
      input_args_dict = (
          (
              arg_name,
              fullargspec.annotations.get(arg_name, None),
              defaults[i],
          )
          for i, arg_name in enumerate(fullargspec.args)
      )
      invalid_types = {}
      for arg_name, arg_type, arg_default in input_args_dict:
          if arg_type is not None and _is_valid_optional_arg(arg_type):
              arg_type = arg_type.__args__[0]
          arg_name = _cli_arg_name(arg_name)
          if arg_type is bool:
              options.append(
                  click.option(
                      f"--{arg_name}",
                      type=click.BOOL,
                      is_flag=True,
                      default=False,
                      required=False,
                  )
              )
          elif arg_type is not None and _is_valid_collection_arg(arg_type):
              member_type = arg_type.__args__[0]
              options.append(
                  click.option(
                      f"--{arg_name}",
                      type=member_type,
                      default=arg_default,
                      required=False,
                      multiple=True,
                  )
              )
          elif arg_type in _ALLOWED_TYPES:
              options.append(
                  click.option(
                      f"--{arg_name}",
                      type=_CLICK_TYPES_MAPPER[arg_type],
                      default=arg_default,
                      required=False if arg_default is not None else True,
                  )
              )
          else:
              invalid_types[arg_name] = arg_type
      if invalid_types:
          raise ValueError(
              f"Invalid argument types: {invalid_types}. CLI functions only "
              f"supports: {_ALLOWED_TYPES} types (including Optional) and "
              f"{_ALLOWED_COLLECTIONS} collections."
          )
      options.append(
          click.command(
              help="Technical wrapper to pass into the `accelerate launch` command."
          )
      )
  
      def wrapper(function: F) -> F:
          for option in reversed(options):
              function = option(function)
          return function
  
      func.__doc__ = (
          f"{func.__doc__}\n\nThis is ZenML-generated CLI wrapper function."
      )
  
      return wrapper(func)
  
  
  @contextmanager
  def create_cli_wrapped_script(
      func: F, flavor: str = "accelerate"
  ) -> Iterator[Tuple[Path, Path]]:
      """Create a script with the CLI-wrapped function.
  
      Args:
          func: The function to use.
          flavor: The flavor to use.
  
      Yields:
          The paths of the script and the output.
  
      Raises:
          ValueError: If the function is not defined in a module.
      """
      try:
          random_name = random_str(20)
          script_path = Path(random_name + ".py")
          output_path = Path(random_name + ".out")
  
          module = inspect.getmodule(func)
          if module is None:
              raise ValueError(
                  f"Function `{func.__name__}` must be defined in a "
                  "module to be used with Accelerate."
              )
  
          with open(script_path, "w") as f:
              if path := module.__file__:
                  func_path = str(Path(path).parent)
                  relative_path = path.replace(func_path, "").lstrip(os.sep)
                  relative_path = os.path.splitext(relative_path)[0]
                  clean_module_name = ".".join(relative_path.split(os.sep))
                  script = _CLI_WRAPPED_SCRIPT_TEMPLATE_HEADER.format(
                      func_path=func_path,
                      func_module=clean_module_name,
                      func_name=func.__name__,
                  )
                  script += _CLI_WRAPPED_MAINS[flavor].format(
                      output_file=str(output_path.absolute())
                  )
                  f.write(script)
              else:
                  raise ValueError(
                      f"Cannot find module file path for function `{func.__name__}`."
                  )
  
          logger.debug(f"Created script:\n\n{script}")
  
          yield script_path, output_path
      finally:
          if script_path.exists():
              script_path.unlink()
          if output_path.exists():
              output_path.unlink()
  