module_utils.py 2.55 KB
Newer Older
1
import importlib.util
2
import warnings
3
from functools import wraps
4
from typing import Optional
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28


def is_module_available(*modules: str) -> bool:
    r"""Returns if a top-level module with :attr:`name` exists *without**
    importing it. This is generally safer than try-catch block around a
    `import X`. It avoids third party libraries breaking assumptions of some of
    our tests, e.g., setting multiprocessing start method when imported
    (see librosa/#747, torchvision/#544).
    """
    return all(importlib.util.find_spec(m) is not None for m in modules)


def requires_module(*modules: str):
    """Decorate function to give error message if invoked without required optional modules.

    This decorator is to give better error message to users rather
    than raising ``NameError:  name 'module' is not defined`` at random places.
    """
    missing = [m for m in modules if not is_module_available(m)]

    if not missing:
        # fall through. If all the modules are available, no need to decorate
        def decorator(func):
            return func
29

30
    else:
31
        req = f"module: {missing[0]}" if len(missing) == 1 else f"modules: {missing}"
32
33
34
35

        def decorator(func):
            @wraps(func)
            def wrapped(*args, **kwargs):
36
37
                raise RuntimeError(f"{func.__module__}.{func.__name__} requires {req}")

38
            return wrapped
39

40
    return decorator
41
42
43
44
45
46


def deprecated(direction: str, version: Optional[str] = None):
    """Decorator to add deprecation message

    Args:
47
48
        direction (str): Migration steps to be given to users.
        version (str or int): The version when the object will be removed
49
    """
50

51
    def decorator(func):
52
53
54
        @wraps(func)
        def wrapped(*args, **kwargs):
            message = (
55
                f"{func.__module__}.{func.__name__} has been deprecated "
56
                f'and will be removed from {"future" if version is None else version} release. '
57
58
                f"{direction}"
            )
59
60
            warnings.warn(message, stacklevel=2)
            return func(*args, **kwargs)
61

62
        return wrapped
63

64
    return decorator
Caroline Chen's avatar
Caroline Chen committed
65
66


67
68
def fail_with_message(message):
    """Generate decorator to give users message about missing TorchAudio extension."""
moto's avatar
moto committed
69

70
71
72
73
    def decorator(func):
        @wraps(func)
        def wrapped(*args, **kwargs):
            raise RuntimeError(f"{func.__module__}.{func.__name__} {message}")
74

75
        return wrapped
76

77
78
79
    return decorator


80
81
82
def no_op(func):
    """Op-op decorator. Used in place of fail_with_message when a functionality that requires extension works fine."""
    return func