content: #  Copyright (c) ZenML GmbH 2020. 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.
  """Logger implementation."""
  
  import logging
  import os
  import re
  import sys
  from contextvars import ContextVar
  from typing import Any, Dict, Optional
  
  from rich.traceback import install as rich_tb_install
  
  from zenml.constants import (
      ENABLE_RICH_TRACEBACK,
      ENV_ZENML_LOGGING_COLORS_DISABLED,
      ENV_ZENML_SUPPRESS_LOGS,
      ZENML_LOGGING_VERBOSITY,
      ZENML_STORAGE_LOGGING_VERBOSITY,
      handle_bool_env_var,
  )
  from zenml.enums import LoggingLevels
  
  ZENML_LOGGING_COLORS_DISABLED = handle_bool_env_var(
      ENV_ZENML_LOGGING_COLORS_DISABLED, False
  )
  
  step_names_in_console: ContextVar[bool] = ContextVar(
      "step_names_in_console", default=False
  )
  
  _original_stdout_write: Optional[Any] = None
  _original_stderr_write: Optional[Any] = None
  _stdout_wrapped: bool = False
  _stderr_wrapped: bool = False
  
  
  class _ZenMLStdoutStream:
      """Stream that writes to the original stdout, bypassing the ZenML wrapper.
  
      This ensures console logging doesn't trigger the LoggingContext wrapper,
      preventing duplicate log entries in stored logs.
      """
  
      def write(self, text: str) -> Any:
          """Write text to the original stdout.
  
          Args:
              text: The text to write.
  
          Returns:
              The number of characters written.
          """
          if _original_stdout_write:
              return _original_stdout_write(text)
          return sys.stdout.write(text)
  
      def flush(self) -> None:
          """Flush the stdout buffer."""
          sys.stdout.flush()
  
  
  def get_logger(logger_name: str) -> logging.Logger:
      """Main function to get logger name,.
  
      Args:
          logger_name: Name of logger to initialize.
  
      Returns:
          A logger object.
      """
      return logging.getLogger(logger_name)
  
  
  def _add_step_name_to_message(message: str) -> str:
      """Adds the step name to the message.
  
      Args:
          message: The message to add the step name to.
  
      Returns:
          The message with the step name added.
      """
      try:
          if step_names_in_console.get():
              from zenml.steps import get_step_context
  
              step_context = get_step_context()
  
              if step_context and message not in ["\n", ""]:
                  # For progress bar updates (with \r), inject the step name after the \r
                  if "\r" in message:
                      message = message.replace(
                          "\r", f"\r[{step_context.step_name}] "
                      )
                  else:
                      message = f"[{step_context.step_name}] {message}"
      except Exception:
          # If we can't get step context, just use the original message
          pass
  
      return message
  
  
  class ConsoleFormatter(logging.Formatter):
      """Formats logs according to custom specifications."""
  
      grey: str = "\x1b[90m"
      white: str = "\x1b[37m"
      pink: str = "\x1b[35m"
      green: str = "\x1b[32m"
      yellow: str = "\x1b[33m"
      red: str = "\x1b[31m"
      cyan: str = "\x1b[1;36m"
      bold_red: str = "\x1b[31;1m"
      purple: str = "\x1b[38;5;105m"
      blue: str = "\x1b[34m"
      reset: str = "\x1b[0m"
  
      def _get_format_template(self, record: logging.LogRecord) -> str:
          """Get the format template based on the logging level.
  
          Args:
              record: The log record to format.
  
          Returns:
              The format template string.
          """
          if get_logging_level() == LoggingLevels.DEBUG:
              return "%(asctime)s - %(name)s - %(levelname)s - %(message)s (%(filename)s:%(lineno)d)"
          else:
              return "%(message)s"
  
      COLORS: Dict[LoggingLevels, str] = {
          LoggingLevels.DEBUG: grey,
          LoggingLevels.INFO: white,
          LoggingLevels.WARN: yellow,
          LoggingLevels.ERROR: red,
          LoggingLevels.CRITICAL: bold_red,
      }
  
      def format(self, record: logging.LogRecord) -> str:
          """Converts a log record to a (colored) string.
  
          Args:
              record: LogRecord generated by the code.
  
          Returns:
              A string formatted according to specifications.
          """
          format_template = self._get_format_template(record)
  
          message = record.getMessage()
          try:
              if step_names_in_console.get():
                  message = _add_step_name_to_message(message)
          except Exception:
              pass
  
          modified_record = logging.LogRecord(
              name=record.name,
              level=record.levelno,
              pathname=record.pathname,
              lineno=record.lineno,
              msg=message,
              args=(),
              exc_info=record.exc_info,
          )
  
          if ZENML_LOGGING_COLORS_DISABLED:
              formatter = logging.Formatter(format_template)
              return formatter.format(modified_record)
          else:
              log_fmt = (
                  self.COLORS[LoggingLevels(record.levelno)]
                  + format_template
                  + self.reset
              )
              formatter = logging.Formatter(log_fmt)
              formatted_message = formatter.format(modified_record)
              quoted_groups = re.findall("`([^`]*)`", formatted_message)
              for quoted in quoted_groups:
                  formatted_message = formatted_message.replace(
                      "`" + quoted + "`",
                      self.reset
                      + self.purple
                      + quoted
                      + self.COLORS.get(LoggingLevels(record.levelno)),
                  )
  
              url_pattern = r"http[s]?://(?:[a-zA-Z]|[0-9]|[$-_@.&+]|[!*\\(\\),]|(?:%[0-9a-fA-F][0-9a-fA-F]))+"
              urls = re.findall(url_pattern, formatted_message)
              for url in urls:
                  formatted_message = formatted_message.replace(
                      url,
                      self.reset
                      + self.blue
                      + url
                      + self.COLORS.get(LoggingLevels(record.levelno)),
                  )
              return formatted_message
  
  
  def get_logging_level() -> LoggingLevels:
      """Get logging level from the env variable.
  
      Returns:
          The logging level.
  
      Raises:
          KeyError: If the logging level is not found.
      """
      verbosity = ZENML_LOGGING_VERBOSITY.upper()
      if verbosity not in LoggingLevels.__members__:
          raise KeyError(
              f"Verbosity must be one of {list(LoggingLevels.__members__.keys())}"
          )
  
      if ZENML_STORAGE_LOGGING_VERBOSITY is not None:
          get_logger(__name__).warning(
              "The ZENML_STORAGE_LOGGING_VERBOSITY is no longer supported. "
              "Please use the ZENML_LOGGING_VERBOSITY instead."
          )
  
      return LoggingLevels[verbosity]
  
  
  def set_root_verbosity() -> None:
      """Set the root verbosity."""
      level = get_logging_level()
      if level != LoggingLevels.NOTSET:
          if ENABLE_RICH_TRACEBACK:
              rich_tb_install(show_locals=(level == LoggingLevels.DEBUG))
  
          logging.root.setLevel(level=level.value)
          get_logger(__name__).debug(
              f"Logging set to level: {logging.getLevelName(level.value)}"
          )
      else:
          logging.disable(sys.maxsize)
          logging.getLogger().disabled = True
          get_logger(__name__).debug("Logging NOTSET")
  
  
  def _wrapped_write(original_write: Any, stream_name: str) -> Any:
      """Wrap stdout/stderr write method to route logs to LoggingContext.
  
      Args:
          original_write: The original write method.
          stream_name: The name of the stream.
  
      Returns:
          The wrapped write method.
      """
  
      def wrapped_write(text: str) -> Any:
          """Write method that routes logs through LoggingContext.
  
          Args:
              text: The text to write.
  
          Returns:
              The result of the original write method.
          """
          from zenml.utils.logging_utils import LoggingContext
  
          level = logging.INFO if stream_name == "stdout" else logging.ERROR
  
          if logging.root.isEnabledFor(level):
              record = logging.LogRecord(
                  name=stream_name,
                  level=level,
                  pathname="",
                  lineno=0,
                  msg=text,
                  args=(),
                  exc_info=None,
                  func="",
              )
              LoggingContext.emit(record)
  
          return original_write(text)
  
      return wrapped_write
  
  
  def wrap_stdout_stderr() -> None:
      """Wrap stdout and stderr write methods to route through LoggingContext."""
      global _stdout_wrapped, _stderr_wrapped
      global _original_stdout_write, _original_stderr_write
  
      if not _stdout_wrapped:
          _original_stdout_write = getattr(sys.stdout, "write")
          setattr(
              sys.stdout,
              "write",
              _wrapped_write(_original_stdout_write, "stdout"),
          )
          _stdout_wrapped = True
  
      if not _stderr_wrapped:
          _original_stderr_write = getattr(sys.stderr, "write")
          setattr(
              sys.stderr,
              "write",
              _wrapped_write(_original_stderr_write, "stderr"),
          )
          _stderr_wrapped = True
  
  
  class ZenMLLoggingHandler(logging.Handler):
      """Custom handler that routes logs through LoggingContext."""
  
      def emit(self, record: logging.LogRecord) -> None:
          """Emit a log record through LoggingContext.
  
          Args:
              record: The log record to emit.
          """
          from zenml.utils.logging_utils import LoggingContext
  
          LoggingContext.emit(record)
  
  
  def get_console_handler() -> logging.Handler:
      """Get console handler that writes to original stdout.
  
      Returns:
          A console handler.
      """
      handler = logging.StreamHandler(_ZenMLStdoutStream())
      handler.setFormatter(ConsoleFormatter())
      return handler
  
  
  def get_zenml_handler() -> logging.Handler:
      """Get ZenML handler that routes logs through LoggingContext.
  
      Returns:
          A ZenML handler.
      """
      return ZenMLLoggingHandler()
  
  
  def init_logging() -> None:
      """Initialize the logging system."""
      set_root_verbosity()
      wrap_stdout_stderr()
  
      # Add both handlers to the root logger
      root_logger = logging.getLogger()
  
      # Console handler - writes to original stdout
      root_logger.addHandler(get_console_handler())
  
      # ZenML handler - routes through LoggingContext
      root_logger.addHandler(get_zenml_handler())
  
      # Mute tensorflow cuda warnings
      os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
  
      # logging capture warnings
      logging.captureWarnings(True)
  
      # Enable logs if environment variable SUPPRESS_ZENML_LOGS is not set to True
      suppress_zenml_logs: bool = handle_bool_env_var(
          ENV_ZENML_SUPPRESS_LOGS, True
      )
      if suppress_zenml_logs:
          # suppress logger info messages
          suppressed_logger_names = [
              "urllib3",
              "azure.core.pipeline.policies.http_logging_policy",
              "grpc",
              "requests",
              "kfp",
              "tensorflow",
          ]
          for logger_name in suppressed_logger_names:
              logging.getLogger(logger_name).setLevel(logging.WARNING)
  
          # disable logger messages
          disabled_logger_names = [
              "rdbms_metadata_access_object",
              "backoff",
              "segment",
          ]
          for logger_name in disabled_logger_names:
              logging.getLogger(logger_name).setLevel(logging.WARNING)
              logging.getLogger(logger_name).disabled = True
  