engine.py

"""Coverage controllers for use by pytest-cov and nose-cov."""

import contextlib
import copy
import functools
import os
import random
import shutil
import socket
import sys
import warnings
from io import StringIO
from pathlib import Path

import coverage
from coverage.data import CoverageData
from coverage.sqldata import filename_suffix

from . import CentralCovContextWarning
from . import DistCovError
from .embed import cleanup


class BrokenCovConfigError(Exception):
    pass


class _NullFile:
    @staticmethod
    def write(v):
        pass


@contextlib.contextmanager
def _backup(obj, attr):
    backup = getattr(obj, attr)
    try:
        setattr(obj, attr, copy.copy(backup))
        yield
    finally:
        setattr(obj, attr, backup)


def _ensure_topdir(meth):
    @functools.wraps(meth)
    def ensure_topdir_wrapper(self, *args, **kwargs):
        try:
            original_cwd = Path.cwd()
        except OSError:
            # Looks like it's gone, this is non-ideal because a side-effect will
            # be introduced in the tests here but we can't do anything about it.
            original_cwd = None
        os.chdir(self.topdir)
        try:
            return meth(self, *args, **kwargs)
        finally:
            if original_cwd is not None:
                os.chdir(original_cwd)

    return ensure_topdir_wrapper


def _data_suffix(name):
    return f'{filename_suffix(True)}.{name}'


class CovController:
    """Base class for different plugin implementations."""

    def __init__(self, cov_source, cov_report, cov_config, cov_append, cov_branch, config=None, nodeid=None):
        """Get some common config used by multiple derived classes."""
        self.cov_source = cov_source
        self.cov_report = cov_report
        self.cov_config = cov_config
        self.cov_append = cov_append
        self.cov_branch = cov_branch
        self.config = config
        self.nodeid = nodeid

        self.cov = None
        self.combining_cov = None
        self.data_file = None
        self.node_descs = set()
        self.failed_workers = []
        self.topdir = os.fspath(Path.cwd())
        self.is_collocated = None

    @contextlib.contextmanager
    def ensure_topdir(self):
        original_cwd = Path.cwd()
        os.chdir(self.topdir)
        yield
        os.chdir(original_cwd)

    @_ensure_topdir
    def pause(self):
        self.cov.stop()
        self.unset_env()

    @_ensure_topdir
    def resume(self):
        self.cov.start()
        self.set_env()

    @_ensure_topdir
    def set_env(self):
        """Put info about coverage into the env so that subprocesses can activate coverage."""
        if self.cov_source is None:
            os.environ['COV_CORE_SOURCE'] = os.pathsep
        else:
            os.environ['COV_CORE_SOURCE'] = os.pathsep.join(
                # We must make the paths absolute to ensure that subprocesses started in different
                # directories still find the same sources: https://github.com/nedbat/coveragepy/issues/1942.
                # Unfortunately, coverage.py doesn't provide a mechanism
                # for unambiguously specifying source_dirs, so we have to detect if a source is
                # a directory or a package by looking in the filesystem: https://github.com/nedbat/coveragepy/issues/1942.
                os.path.abspath(source)  # noqa: PTH100 #  see note below for why we're using `os.path.abspath`
                if Path(source).is_dir()
                else source
                for source in self.cov_source
            )
        config_file = Path(self.cov_config)
        if config_file.exists():
            os.environ['COV_CORE_CONFIG'] = os.fspath(config_file.resolve())
        else:
            os.environ['COV_CORE_CONFIG'] = os.pathsep
        # this still uses the old abspath cause apparently Python 3.9 on Windows has a buggy Path.resolve()
        os.environ['COV_CORE_DATAFILE'] = os.path.abspath(self.cov.config.data_file)  # noqa: PTH100
        if self.cov_branch:
            os.environ['COV_CORE_BRANCH'] = 'enabled'

    @staticmethod
    def unset_env():
        """Remove coverage info from env."""
        os.environ.pop('COV_CORE_SOURCE', None)
        os.environ.pop('COV_CORE_CONFIG', None)
        os.environ.pop('COV_CORE_DATAFILE', None)
        os.environ.pop('COV_CORE_BRANCH', None)
        os.environ.pop('COV_CORE_CONTEXT', None)

    @staticmethod
    def get_node_desc(platform, version_info):
        """Return a description of this node."""

        return 'platform {}, python {}'.format(platform, '{}.{}.{}-{}-{}'.format(*version_info[:5]))

    @staticmethod
    def get_width():
        # taken from https://github.com/pytest-dev/pytest/blob/33c7b05a/src/_pytest/_io/terminalwriter.py#L26
        width, _ = shutil.get_terminal_size(fallback=(80, 24))
        # The Windows get_terminal_size may be bogus, let's sanify a bit.
        if width < 40:
            width = 80
        return width

    def sep(self, stream, s, txt):
        if hasattr(stream, 'sep'):
            stream.sep(s, txt)
        else:
            fullwidth = self.get_width()
            # taken from https://github.com/pytest-dev/pytest/blob/33c7b05a/src/_pytest/_io/terminalwriter.py#L126
            # The goal is to have the line be as long as possible
            # under the condition that len(line) <= fullwidth.
            if sys.platform == 'win32':
                # If we print in the last column on windows we are on a
                # new line but there is no way to verify/neutralize this
                # (we may not know the exact line width).
                # So let's be defensive to avoid empty lines in the output.
                fullwidth -= 1
            N = max((fullwidth - len(txt) - 2) // (2 * len(s)), 1)
            fill = s * N
            line = f'{fill} {txt} {fill}'
            # In some situations there is room for an extra sepchar at the right,
            # in particular if we consider that with a sepchar like "_ " the
            # trailing space is not important at the end of the line.
            if len(line) + len(s.rstrip()) <= fullwidth:
                line += s.rstrip()
            # (end of terminalwriter borrowed code)
            line += '\n\n'
            stream.write(line)

    @_ensure_topdir
    def summary(self, stream):
        """Produce coverage reports."""
        total = None

        if not self.cov_report:
            with _backup(self.cov, 'config'):
                return self.cov.report(show_missing=True, ignore_errors=True, file=_NullFile)

        # Output coverage section header.
        if len(self.node_descs) == 1:
            self.sep(stream, '_', f'coverage: {"".join(self.node_descs)}')
        else:
            self.sep(stream, '_', 'coverage')
            for node_desc in sorted(self.node_descs):
                self.sep(stream, ' ', f'{node_desc}')

        # Report on any failed workers.
        if self.failed_workers:
            self.sep(stream, '_', 'coverage: failed workers')
            stream.write('The following workers failed to return coverage data, ensure that pytest-cov is installed on these workers.\n')
            for node in self.failed_workers:
                stream.write(f'{node.gateway.id}\n')

        # Produce terminal report if wanted.
        if any(x in self.cov_report for x in ['term', 'term-missing']):
            options = {
                'show_missing': ('term-missing' in self.cov_report) or None,
                'ignore_errors': True,
                'file': stream,
            }
            skip_covered = isinstance(self.cov_report, dict) and 'skip-covered' in self.cov_report.values()
            options.update({'skip_covered': skip_covered or None})
            with _backup(self.cov, 'config'):
                total = self.cov.report(**options)

        # Produce annotated source code report if wanted.
        if 'annotate' in self.cov_report:
            annotate_dir = self.cov_report['annotate']

            with _backup(self.cov, 'config'):
                self.cov.annotate(ignore_errors=True, directory=annotate_dir)
            # We need to call Coverage.report here, just to get the total
            # Coverage.annotate don't return any total and we need it for --cov-fail-under.

            with _backup(self.cov, 'config'):
                total = self.cov.report(ignore_errors=True, file=_NullFile)
            if annotate_dir:
                stream.write(f'Coverage annotated source written to dir {annotate_dir}\n')
            else:
                stream.write('Coverage annotated source written next to source\n')

        # Produce html report if wanted.
        if 'html' in self.cov_report:
            output = self.cov_report['html']
            with _backup(self.cov, 'config'):
                total = self.cov.html_report(ignore_errors=True, directory=output)
            stream.write(f'Coverage HTML written to dir {self.cov.config.html_dir if output is None else output}\n')

        # Produce xml report if wanted.
        if 'xml' in self.cov_report:
            output = self.cov_report['xml']
            with _backup(self.cov, 'config'):
                total = self.cov.xml_report(ignore_errors=True, outfile=output)
            stream.write(f'Coverage XML written to file {self.cov.config.xml_output if output is None else output}\n')

        # Produce json report if wanted
        if 'json' in self.cov_report:
            output = self.cov_report['json']
            with _backup(self.cov, 'config'):
                total = self.cov.json_report(ignore_errors=True, outfile=output)
            stream.write('Coverage JSON written to file %s\n' % (self.cov.config.json_output if output is None else output))

        # Produce lcov report if wanted.
        if 'lcov' in self.cov_report:
            output = self.cov_report['lcov']
            with _backup(self.cov, 'config'):
                self.cov.lcov_report(ignore_errors=True, outfile=output)

                # We need to call Coverage.report here, just to get the total
                # Coverage.lcov_report doesn't return any total and we need it for --cov-fail-under.
                total = self.cov.report(ignore_errors=True, file=_NullFile)

            stream.write(f'Coverage LCOV written to file {self.cov.config.lcov_output if output is None else output}\n')

        return total


class Central(CovController):
    """Implementation for centralised operation."""

    @_ensure_topdir
    def start(self):
        cleanup()

        self.cov = coverage.Coverage(
            source=self.cov_source,
            branch=self.cov_branch,
            data_suffix=_data_suffix('c'),
            config_file=self.cov_config,
        )
        if self.cov.config.dynamic_context == 'test_function':
            message = (
                'Detected dynamic_context=test_function in coverage configuration. '
                'This is unnecessary as this plugin provides the more complete --cov-context option.'
            )
            warnings.warn(CentralCovContextWarning(message), stacklevel=1)

        self.combining_cov = coverage.Coverage(
            source=self.cov_source,
            branch=self.cov_branch,
            data_suffix=_data_suffix('cc'),
            data_file=os.path.abspath(self.cov.config.data_file),  # noqa: PTH100
            config_file=self.cov_config,
        )

        # Erase or load any previous coverage data and start coverage.
        if not self.cov_append:
            self.cov.erase()
        self.cov.start()
        self.set_env()

    @_ensure_topdir
    def finish(self):
        """Stop coverage, save data to file and set the list of coverage objects to report on."""

        self.unset_env()
        self.cov.stop()
        self.cov.save()

        self.cov = self.combining_cov
        self.cov.load()
        self.cov.combine()
        self.cov.save()

        node_desc = self.get_node_desc(sys.platform, sys.version_info)
        self.node_descs.add(node_desc)


class DistMaster(CovController):
    """Implementation for distributed master."""

    @_ensure_topdir
    def start(self):
        cleanup()

        self.cov = coverage.Coverage(
            source=self.cov_source,
            branch=self.cov_branch,
            data_suffix=_data_suffix('m'),
            config_file=self.cov_config,
        )
        if self.cov.config.dynamic_context == 'test_function':
            raise DistCovError(
                'Detected dynamic_context=test_function in coverage configuration. '
                'This is known to cause issues when using xdist, see: https://github.com/pytest-dev/pytest-cov/issues/604\n'
                'It is recommended to use --cov-context instead.'
            )
        self.cov._warn_no_data = False
        self.cov._warn_unimported_source = False
        self.cov._warn_preimported_source = False
        self.combining_cov = coverage.Coverage(
            source=self.cov_source,
            branch=self.cov_branch,
            data_suffix=_data_suffix('mc'),
            data_file=os.path.abspath(self.cov.config.data_file),  # noqa: PTH100
            config_file=self.cov_config,
        )
        if not self.cov_append:
            self.cov.erase()
        self.cov.start()
        self.cov.config.paths['source'] = [self.topdir]

    def configure_node(self, node):
        """Workers need to know if they are collocated and what files have moved."""

        node.workerinput.update(
            {
                'cov_master_host': socket.gethostname(),
                'cov_master_topdir': self.topdir,
                'cov_master_rsync_roots': [str(root) for root in node.nodemanager.roots],
            }
        )

    def testnodedown(self, node, error):
        """Collect data file name from worker."""

        # If worker doesn't return any data then it is likely that this
        # plugin didn't get activated on the worker side.
        output = getattr(node, 'workeroutput', {})
        if 'cov_worker_node_id' not in output:
            self.failed_workers.append(node)
            return

        # If worker is not collocated then we must save the data file
        # that it returns to us.
        if 'cov_worker_data' in output:
            data_suffix = '%s.%s.%06d.%s' % (  # noqa: UP031
                socket.gethostname(),
                os.getpid(),
                random.randint(0, 999999),  # noqa: S311
                output['cov_worker_node_id'],
            )

            cov = coverage.Coverage(source=self.cov_source, branch=self.cov_branch, data_suffix=data_suffix, config_file=self.cov_config)
            cov.start()
            if coverage.version_info < (5, 0):
                data = CoverageData()
                data.read_fileobj(StringIO(output['cov_worker_data']))
                cov.data.update(data)
            else:
                data = CoverageData(no_disk=True, suffix='should-not-exist')
                data.loads(output['cov_worker_data'])
                cov.get_data().update(data)
            cov.stop()
            cov.save()
            path = output['cov_worker_path']
            self.cov.config.paths['source'].append(path)

        # Record the worker types that contribute to the data file.
        rinfo = node.gateway._rinfo()
        node_desc = self.get_node_desc(rinfo.platform, rinfo.version_info)
        self.node_descs.add(node_desc)

    @_ensure_topdir
    def finish(self):
        """Combines coverage data and sets the list of coverage objects to report on."""

        # Combine all the suffix files into the data file.
        self.cov.stop()
        self.cov.save()
        self.cov = self.combining_cov
        self.cov.load()
        self.cov.combine()
        self.cov.save()


class DistWorker(CovController):
    """Implementation for distributed workers."""

    @_ensure_topdir
    def start(self):
        cleanup()

        # Determine whether we are collocated with master.
        self.is_collocated = (
            socket.gethostname() == self.config.workerinput['cov_master_host']
            and self.topdir == self.config.workerinput['cov_master_topdir']
        )

        # If we are not collocated then rewrite master paths to worker paths.
        if not self.is_collocated:
            master_topdir = self.config.workerinput['cov_master_topdir']
            worker_topdir = self.topdir
            if self.cov_source is not None:
                self.cov_source = [source.replace(master_topdir, worker_topdir) for source in self.cov_source]
            self.cov_config = self.cov_config.replace(master_topdir, worker_topdir)

        # Erase any previous data and start coverage.
        self.cov = coverage.Coverage(
            source=self.cov_source,
            branch=self.cov_branch,
            data_suffix=_data_suffix(f'w{self.nodeid}'),
            config_file=self.cov_config,
        )
        self.cov.start()
        self.set_env()

    @_ensure_topdir
    def finish(self):
        """Stop coverage and send relevant info back to the master."""
        self.unset_env()
        self.cov.stop()

        if self.is_collocated:
            # We don't combine data if we're collocated - we can get
            # race conditions in the .combine() call (it's not atomic)
            # The data is going to be combined in the master.
            self.cov.save()

            # If we are collocated then just inform the master of our
            # data file to indicate that we have finished.
            self.config.workeroutput['cov_worker_node_id'] = self.nodeid
        else:
            self.cov.combine()
            self.cov.save()
            # If we are not collocated then add the current path
            # and coverage data to the output so we can combine
            # it on the master node.

            # Send all the data to the master over the channel.
            if coverage.version_info < (5, 0):
                buff = StringIO()
                self.cov.data.write_fileobj(buff)
                data = buff.getvalue()
            else:
                data = self.cov.get_data().dumps()

            self.config.workeroutput.update(
                {
                    'cov_worker_path': self.topdir,
                    'cov_worker_node_id': self.nodeid,
                    'cov_worker_data': data,
                }
            )

    def summary(self, stream):
        """Only the master reports so do nothing."""