1# vim:fileencoding=utf-8:noet
2from __future__ import (unicode_literals, division, absolute_import, print_function)
3
4import os
5
6from multiprocessing import cpu_count as _cpu_count
7
8from powerline.lib.threaded import ThreadedSegment
9from powerline.lib import add_divider_highlight_group
10from powerline.segments import with_docstring
11
12
13cpu_count = None
14
15
16def system_load(pl, format='{avg:.1f}', threshold_good=1, threshold_bad=2,
17                track_cpu_count=False, short=False):
18	'''Return system load average.
19
20	Highlights using ``system_load_good``, ``system_load_bad`` and
21	``system_load_ugly`` highlighting groups, depending on the thresholds
22	passed to the function.
23
24	:param str format:
25		format string, receives ``avg`` as an argument
26	:param float threshold_good:
27		threshold for gradient level 0: any normalized load average below this
28		value will have this gradient level.
29	:param float threshold_bad:
30		threshold for gradient level 100: any normalized load average above this
31		value will have this gradient level. Load averages between
32		``threshold_good`` and ``threshold_bad`` receive gradient level that
33		indicates relative position in this interval:
34		(``100 * (cur-good) / (bad-good)``).
35		Note: both parameters are checked against normalized load averages.
36	:param bool track_cpu_count:
37		if True powerline will continuously poll the system to detect changes
38		in the number of CPUs.
39	:param bool short:
40		if True only the sys load over last 1 minute will be displayed.
41
42	Divider highlight group used: ``background:divider``.
43
44	Highlight groups used: ``system_load_gradient`` (gradient) or ``system_load``.
45	'''
46	global cpu_count
47	try:
48		cpu_num = cpu_count = _cpu_count() if cpu_count is None or track_cpu_count else cpu_count
49	except NotImplementedError:
50		pl.warn('Unable to get CPU count: method is not implemented')
51		return None
52	ret = []
53	for avg in os.getloadavg():
54		normalized = avg / cpu_num
55		if normalized < threshold_good:
56			gradient_level = 0
57		elif normalized < threshold_bad:
58			gradient_level = (normalized - threshold_good) * 100.0 / (threshold_bad - threshold_good)
59		else:
60			gradient_level = 100
61		ret.append({
62			'contents': format.format(avg=avg),
63			'highlight_groups': ['system_load_gradient', 'system_load'],
64			'divider_highlight_group': 'background:divider',
65			'gradient_level': gradient_level,
66		})
67
68		if short:
69		    return ret
70
71	ret[0]['contents'] += ' '
72	ret[1]['contents'] += ' '
73	return ret
74
75
76try:
77	import psutil
78
79	class CPULoadPercentSegment(ThreadedSegment):
80		interval = 1
81
82		def update(self, old_cpu):
83			return psutil.cpu_percent(interval=None)
84
85		def run(self):
86			while not self.shutdown_event.is_set():
87				try:
88					self.update_value = psutil.cpu_percent(interval=self.interval)
89				except Exception as e:
90					self.exception('Exception while calculating cpu_percent: {0}', str(e))
91
92		def render(self, cpu_percent, format='{0:.0f}%', **kwargs):
93			return [{
94				'contents': format.format(cpu_percent),
95				'gradient_level': cpu_percent,
96				'highlight_groups': ['cpu_load_percent_gradient', 'cpu_load_percent'],
97			}]
98except ImportError:
99	class CPULoadPercentSegment(ThreadedSegment):
100		interval = 1
101
102		@staticmethod
103		def startup(**kwargs):
104			pass
105
106		@staticmethod
107		def start():
108			pass
109
110		@staticmethod
111		def shutdown():
112			pass
113
114		@staticmethod
115		def render(cpu_percent, pl, format='{0:.0f}%', **kwargs):
116			pl.warn('Module “psutil” is not installed, thus CPU load is not available')
117			return None
118
119
120cpu_load_percent = with_docstring(CPULoadPercentSegment(),
121'''Return the average CPU load as a percentage.
122
123Requires the ``psutil`` module.
124
125:param str format:
126	Output format. Accepts measured CPU load as the first argument.
127
128Highlight groups used: ``cpu_load_percent_gradient`` (gradient) or ``cpu_load_percent``.
129''')
130
131
132if os.path.exists('/proc/uptime'):
133	def _get_uptime():
134		with open('/proc/uptime', 'r') as f:
135			return int(float(f.readline().split()[0]))
136elif 'psutil' in globals():
137	from time import time
138
139	if hasattr(psutil, 'boot_time'):
140		def _get_uptime():
141			return int(time() - psutil.boot_time())
142	else:
143		def _get_uptime():
144			return int(time() - psutil.BOOT_TIME)
145else:
146	def _get_uptime():
147		raise NotImplementedError
148
149
150@add_divider_highlight_group('background:divider')
151def uptime(pl, days_format='{days:d}d', hours_format=' {hours:d}h', minutes_format=' {minutes:02d}m',
152		seconds_format=' {seconds:02d}s', shorten_len=3):
153	'''Return system uptime.
154
155	:param str days_format:
156		day format string, will be passed ``days`` as the argument
157	:param str hours_format:
158		hour format string, will be passed ``hours`` as the argument
159	:param str minutes_format:
160		minute format string, will be passed ``minutes`` as the argument
161	:param str seconds_format:
162		second format string, will be passed ``seconds`` as the argument
163	:param int shorten_len:
164		shorten the amount of units (days, hours, etc.) displayed
165
166	Divider highlight group used: ``background:divider``.
167	'''
168	try:
169		seconds = _get_uptime()
170	except NotImplementedError:
171		pl.warn('Unable to get uptime. You should install psutil module')
172		return None
173	minutes, seconds = divmod(seconds, 60)
174	hours, minutes = divmod(minutes, 60)
175	days, hours = divmod(hours, 24)
176	time_formatted = list(filter(None, [
177		days_format.format(days=days) if days_format else None,
178		hours_format.format(hours=hours) if hours_format else None,
179		minutes_format.format(minutes=minutes) if minutes_format else None,
180		seconds_format.format(seconds=seconds) if seconds_format else None,
181	]))
182	first_non_zero = next((i for i, x in enumerate([days, hours, minutes, seconds]) if x != 0))
183	time_formatted = time_formatted[first_non_zero:first_non_zero + shorten_len]
184	return ''.join(time_formatted).strip()
185