local_agent/logger.py

# Copyright 2022 Google LLC
#
# 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.

"""Module for Local Agent logger."""
import logging
from logging import handlers
from typing import Optional


log_handler_map = {}
logger = logging.getLogger(name='local_agent')

_DEFAULT_LOG_FILE = 'local_agent.log'

# Example:
# 20210725 18:23:09.509 W tmp.py:23 This is the log message.
_LOG_FORMAT_PARTS = ('%(asctime)s.%(msecs)03d',  # Time.
                     '%(levelname).1s',  # Level.
                     '%(filename)s:%(lineno)d',  # Where.
                     '%(message)s')
_LOG_FORMAT = ' '.join(_LOG_FORMAT_PARTS)
_DATE_FORMAT = '%Y%m%d %X'


def create_handler(log_file: Optional[str] = None) -> logging.Handler:
    """Creates the logging handler for a specific log file.

    If log_file arugment is provided, we create RotatingFileHandler; otherwise
    we create a StreamHandler.

    Args:
        log_file (str): Path to the log file.

    Returns:
        The logging handler created.
    """
    if log_file is not None:
        handler = handlers.RotatingFileHandler(log_file,
                                               maxBytes=100000,
                                               backupCount=10)
    else:
        handler = logging.StreamHandler()
    handler.setLevel(logging.DEBUG)
    handler.setFormatter(logging.Formatter(_LOG_FORMAT, datefmt=_DATE_FORMAT))
    return handler


def add_file_handler(log_file: str) -> None:
    """Adds a handler to the global logger for logging to the given file.

    We use a global dict to prevent adding multiple file handlers for one
    single file.

    Args:
        log_file: Path of the file to add handler for.
    """
    if log_file in log_handler_map:
        return
    handler = create_handler(log_file)
    log_handler_map[log_file] = handler
    logger.addHandler(handler)


def remove_file_handler(log_file: str) -> None:
    """Removes the logging handler of a specific file.

    If there isn't a handler for the file, no-op.

    Args:
        log_file: Path of the log file.
    """
    if log_file not in log_handler_map:
        return
    handler = log_handler_map[log_file]
    logger.removeHandler(handler)
    del log_handler_map[log_file]


def setup_logger() -> None:
    """Sets up the logger for logging in file and console."""
    logger.setLevel(logging.DEBUG)

    file_handler = create_handler(_DEFAULT_LOG_FILE)
    stream_handler = create_handler()

    logger.addHandler(file_handler)
    logger.addHandler(stream_handler)


def get_logger() -> logging.Logger:
    """Gets the local agent logger for logging.

    This get function helps modules in the local agent package get the local
    agent logger without having to directly importing the "logger" member.

    Returns:
        The local agent logger.
    """
    return logger


setup_logger()