Repository: Lightning-AI/pytorch-lightning
File: tests/helpers/skipif.py (Python)

Code before:
from distutils.version import LooseVersion
from typing import Optional

import pytest
import torch
from pkg_resources import get_distribution

from pytorch_lightning.utilities import _TORCH_QUANTIZE_AVAILABLE

_MISS_QUANT_DEFAULT = 'fbgemm' not in torch.backends.quantized.supported_engines


def create_skipif(
    min_gpus: int = 0,
    min_torch: Optional[str] = None,
    pt_quant: bool = False,
) -> dict:
    """ Creating aggregated arguments for standatd pytest skipif, sot the usecase is::

        @pytest.mark.skipif(**create_skipif(min_torch="99"))
        def test_any_func(...):
            ...

    >>> from pprint import pprint
    >>> pprint(create_skipif(min_torch="99"))
    {'condition': True, 'reason': 'test requires minimal version `torch>=99'}
    >>> pprint(create_skipif(min_torch="0.0"))
    {'condition': False, 'reason': 'Conditions satisfied, going ahead with the test.'}
    """
    conditions = []
    reasons = []

    if min_gpus:
        conditions.append(torch.cuda.device_count() < min_gpus)
        reasons.append(f"multi-GPU machine with at least {min_gpus}")

    if min_torch:
        torch_version = LooseVersion(get_distribution("torch").version)
        conditions.append(torch_version < LooseVersion(min_torch))
        reasons.append(f"minimal version `torch>={min_torch}")

    if pt_quant:
        conditions.append(not _TORCH_QUANTIZE_AVAILABLE or _MISS_QUANT_DEFAULT)
        reasons.append("PyTorch quantization")

    if not any(conditions):
        return dict(condition=False, reason="Conditions satisfied, going ahead with the test.")

    reasons = [rs for cond, rs in zip(conditions, reasons) if cond]
    reason = "test requires " + ' + '.join(reasons)
    return dict(condition=any(conditions), reason=reason)

Review comment (bug):
doctest is failing here, should probably be:
```suggestion
    >>> pprint(skipif_args(min_torch="0.0"))  # doctest: +NORMALIZE_WHITESPACE
    {'condition': False, 'reason': # doctest: +NORMALIZE_WHITESPACE
     'Conditions satisfied, going ahead with the test.'}
```

Code after:
from distutils.version import LooseVersion
from typing import Optional

import pytest
import torch
from pkg_resources import get_distribution

from pytorch_lightning.utilities import _TORCH_QUANTIZE_AVAILABLE


def skipif_args(
    min_gpus: int = 0,
    min_torch: Optional[str] = None,
    quant_available: bool = False,
) -> dict:
    """ Creating aggregated arguments for standard pytest skipif, sot the usecase is::

        @pytest.mark.skipif(**create_skipif(min_torch="99"))
        def test_any_func(...):
            ...

    >>> from pprint import pprint
    >>> pprint(skipif_args(min_torch="99", min_gpus=0))
    {'condition': True, 'reason': 'Required: [torch>=99]'}
    >>> pprint(skipif_args(min_torch="0.0", min_gpus=0))  # doctest: +NORMALIZE_WHITESPACE
    {'condition': False, 'reason': 'Conditions satisfied, going ahead with the test.'}
    """
    conditions = []
    reasons = []

    if min_gpus:
        conditions.append(torch.cuda.device_count() < min_gpus)
        reasons.append(f"GPUs>={min_gpus}")

    if min_torch:
        torch_version = LooseVersion(get_distribution("torch").version)
        conditions.append(torch_version < LooseVersion(min_torch))
        reasons.append(f"torch>={min_torch}")

    if quant_available:
        _miss_default = 'fbgemm' not in torch.backends.quantized.supported_engines
        conditions.append(not _TORCH_QUANTIZE_AVAILABLE or _miss_default)
        reasons.append("PyTorch quantization is available")

    if not any(conditions):
        return dict(condition=False, reason="Conditions satisfied, going ahead with the test.")

    reasons = [rs for cond, rs in zip(conditions, reasons) if cond]
    return dict(condition=any(conditions), reason=f"Required: [{' + '.join(reasons)}]",)