threadedlogger.py

#  Copyright 2014-2015 Nokia Networks
#  Copyright 2016- Robot Framework Foundation
#
#  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.

from __future__ import print_function, with_statement
__version__ = '0.1'

import os
import threading
from contextlib import contextmanager
from robot.api import logger
from robot.utils import get_timestamp


class BaseLogger(object):
    """Base class for custom loggers with same api as ``robot.api.logger``."""

    def trace(self, msg, html=False):
        self.write(msg, 'TRACE', html)

    def debug(self, msg, html=False):
        self.write(msg, 'DEBUG', html)

    def error(self, msg, html=False):
        self.write(msg, 'ERROR', html)

    def info(self, msg, html=False, also_to_console=False):
        self.write(msg, 'INFO', html)
        if also_to_console:
            self.console(msg)

    def warn(self, msg, html=False):
        self.write(msg, 'WARN', html)

    def console(self, msg, newline=True, stream='stdout'):
        logger.console(msg, newline, stream)

    def write(self, msg, level, html=False):
        raise NotImplementedError


class ThreadedLogger(BaseLogger):
    ROBOT_LISTENER_API_VERSION = 2

    def __init__(self, filename="threaded.log"):
        self._filename = filename
        self._thread_counts = {}
        if os.path.exists(self._filename):
            os.remove(filename)
    
    def start_suite(self, name, attributes):
        self._write("--- START SUITE '%s' ---" % name, "INFO")

    def end_suite(self, name, attributes):
        self._write("--- END SUITE '%s' ---" % name, "INFO")

    def start_test(self, name, attributes):
        self._write("--- START TEST '%s' ---" % name, "INFO")

    def end_test(self, name, attributes):
        self._write("--- END TEST '%s' ---" % name, "INFO")
    
    def start_keyword(self, name, attributes):
        self._write("--- START KEYWORD '%s' ---" % name, "INFO")

    def end_keyword(self, name, attributes):
        self._write("--- END KEYWORD '%s' ---" % name, "INFO")
    
    def log_message(self, message):
        if self._thread_counts.get(threading.current_thread(), 0) == 0:
            self._write(message["message"], message["level"])

    def _write(self, msg, level):
        with open(self._filename, "a") as threadedlog:
            threadedlog.write("%s : %s : %s\n" % (get_timestamp(), level, msg))

    def write(self, msg, level, html=False):
        self._write(msg, level)
        with self._thread_count():
            logger.write(msg, level, html)
    
    @contextmanager
    def _thread_count(self):
        current = threading.current_thread()
        if current not in self._thread_counts:
            self._thread_counts[current] = 0
        self._thread_counts[current] += 1
        yield
        self._thread_counts[current] -= 1


LOGGER = ThreadedLogger()
# THIS NEEDS TO BE ALSO ADDED AS A ROBOT_LIBRARY_LISTENER
# FOR ONE GLOBAL LIBRARY!!!!

_CALLED_LIBRARY_LOGGER = False
def libraryLoggerOrStub():
    global _CALLED_LIBRARY_LOGGER
    if _CALLED_LIBRARY_LOGGER:
        stub = lambda: 0
        stub.ROBOT_LISTENER_API_VERSION = 2
        return stub
    _CALLED_LIBRARY_LOGGER = True
    return LOGGER