1"""
2Version agnostic psutil hack to fully support both old (<2.0) and new (>=2.0)
3psutil versions.
4
5The old <1.0 psutil API is dropped in psutil 3.0
6
7Should be removed once support for psutil <2.0 is dropped. (eg RHEL 6)
8
9Built off of http://grodola.blogspot.com/2014/01/psutil-20-porting.html
10"""
11
12
13# No exception handling, as we want ImportError if psutil doesn't exist
14import psutil  # pylint: disable=3rd-party-module-not-gated
15
16if psutil.version_info >= (2, 0):
17    from psutil import *  # pylint: disable=wildcard-import,unused-wildcard-import,3rd-party-module-not-gated
18else:
19    # Import hack to work around bugs in old psutil's
20    # Psuedo "from psutil import *"
21    _globals = globals()
22    for attr in psutil.__all__:
23        _temp = __import__("psutil", globals(), locals(), [attr], 0)
24        try:
25            _globals[attr] = getattr(_temp, attr)
26        except AttributeError:
27            pass
28
29    # Import functions not in __all__
30    # pylint: disable=unused-import,3rd-party-module-not-gated
31    from psutil import disk_partitions
32    from psutil import disk_usage
33
34    # pylint: enable=unused-import,3rd-party-module-not-gated
35
36    # Alias new module functions
37    def boot_time():
38        return psutil.BOOT_TIME
39
40    def cpu_count():
41        return psutil.NUM_CPUS
42
43    # Alias renamed module functions
44    pids = psutil.get_pid_list
45    try:
46        users = psutil.get_users
47    except AttributeError:
48        users = lambda: (_ for _ in ()).throw(
49            NotImplementedError("Your psutil version is too old")
50        )
51
52    # Deprecated in 1.0.1, but not mentioned in blog post
53    if psutil.version_info < (1, 0, 1):
54        net_io_counters = psutil.network_io_counters()
55
56    class Process(psutil.Process):  # pylint: disable=no-init
57        # Reimplement overloaded getters/setters
58        # pylint: disable=arguments-differ
59        def cpu_affinity(self, *args, **kwargs):
60            if args or kwargs:
61                return self.set_cpu_affinity(*args, **kwargs)
62            else:
63                return self.get_cpu_affinity()
64
65        def ionice(self, *args, **kwargs):
66            if args or kwargs:
67                return self.set_ionice(*args, **kwargs)
68            else:
69                return self.get_ionice()
70
71        def nice(self, *args, **kwargs):
72            if args or kwargs:
73                return self.set_nice(*args, **kwargs)
74            else:
75                return self.get_nice()
76
77        def rlimit(self, *args, **kwargs):
78            """
79            set_rlimit and get_limit were not introduced until psutil v1.1.0
80            """
81            if psutil.version_info >= (1, 1, 0):
82                if args or kwargs:
83                    return self.set_rlimit(*args, **kwargs)
84                else:
85                    return self.get_rlimit()
86            else:
87                pass
88
89        # pylint: enable=arguments-differ
90
91    # Alias renamed Process functions
92    _PROCESS_FUNCTION_MAP = {
93        "children": "get_children",
94        "connections": "get_connections",
95        "cpu_percent": "get_cpu_percent",
96        "cpu_times": "get_cpu_times",
97        "io_counters": "get_io_counters",
98        "memory_info": "get_memory_info",
99        "memory_info_ex": "get_ext_memory_info",
100        "memory_maps": "get_memory_maps",
101        "memory_percent": "get_memory_percent",
102        "num_ctx_switches": "get_num_ctx_switches",
103        "num_fds": "get_num_fds",
104        "num_threads": "get_num_threads",
105        "open_files": "get_open_files",
106        "threads": "get_threads",
107        "cwd": "getcwd",
108    }
109
110    for new, old in _PROCESS_FUNCTION_MAP.items():
111        try:
112            setattr(Process, new, psutil.Process.__dict__[old])
113        except KeyError:
114            pass
115