summaryrefslogtreecommitdiffstats
path: root/powerline/segments/common/env.py
blob: bbfe3e2506c80ab4b700305aaac90bc3d36628b8 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
# vim:fileencoding=utf-8:noet
from __future__ import (unicode_literals, division, absolute_import, print_function)

import os

from powerline.lib.unicode import out_u
from powerline.theme import requires_segment_info
from powerline.segments import Segment, with_docstring


@requires_segment_info
def environment(pl, segment_info, variable=None):
	'''Return the value of any defined environment variable

	:param string variable:
		The environment variable to return if found
	'''
	return segment_info['environ'].get(variable, None)


@requires_segment_info
def virtualenv(pl, segment_info, ignore_venv=False, ignore_conda=False, ignored_names=("venv", ".venv")):
	'''Return the name of the current Python or conda virtualenv.
	:param list ignored_names:
		Names of venvs to ignore. Will then get the name of the venv by ascending to the parent directory
	:param bool ignore_venv:
		Whether to ignore virtual environments. Default is False.
	:param bool ignore_conda:
		Whether to ignore conda environments. Default is False.
	'''
	if not ignore_venv:
		for candidate in reversed(segment_info['environ'].get('VIRTUAL_ENV', '').split("/")):
			if candidate and candidate not in ignored_names:
				return candidate
	if not ignore_conda:
		for candidate in reversed(segment_info['environ'].get('CONDA_DEFAULT_ENV', '').split("/")):
			if candidate and candidate not in ignored_names:
				return candidate
	return None


@requires_segment_info
class CwdSegment(Segment):
	def argspecobjs(self):
		for obj in super(CwdSegment, self).argspecobjs():
			yield obj
		yield 'get_shortened_path', self.get_shortened_path

	def omitted_args(self, name, method):
		if method is self.get_shortened_path:
			return ()
		else:
			return super(CwdSegment, self).omitted_args(name, method)

	def get_shortened_path(self, pl, segment_info, shorten_home=True, **kwargs):
		try:
			path = out_u(segment_info['getcwd']())
		except OSError as e:
			if e.errno == 2:
				# user most probably deleted the directory
				# this happens when removing files from Mercurial repos for example
				pl.warn('Current directory not found')
				return '[not found]'
			else:
				raise
		if shorten_home:
			home = segment_info['home']
			if home:
				home = out_u(home)
				if path.startswith(home):
					path = '~' + path[len(home):]
		return path

	def __call__(self, pl, segment_info,
	             dir_shorten_len=None,
	             dir_limit_depth=None,
	             use_path_separator=False,
	             ellipsis='...',
	             **kwargs):
		cwd = self.get_shortened_path(pl, segment_info, **kwargs)
		cwd_split = cwd.split(os.sep)
		cwd_split_len = len(cwd_split)
		cwd = [i[0:dir_shorten_len] if dir_shorten_len and i else i for i in cwd_split[:-1]] + [cwd_split[-1]]
		if dir_limit_depth and cwd_split_len > dir_limit_depth + 1:
			del(cwd[0:-dir_limit_depth])
			if ellipsis is not None:
				cwd.insert(0, ellipsis)
		ret = []
		if not cwd[0]:
			cwd[0] = '/'
		draw_inner_divider = not use_path_separator
		for part in cwd:
			if not part:
				continue
			if use_path_separator:
				part += os.sep
			ret.append({
				'contents': part,
				'divider_highlight_group': 'cwd:divider',
				'draw_inner_divider': draw_inner_divider,
			})
		ret[-1]['highlight_groups'] = ['cwd:current_folder', 'cwd']
		if use_path_separator:
			ret[-1]['contents'] = ret[-1]['contents'][:-1]
			if len(ret) > 1 and ret[0]['contents'][0] == os.sep:
				ret[0]['contents'] = ret[0]['contents'][1:]
		return ret


cwd = with_docstring(CwdSegment(),
'''Return the current working directory.

Returns a segment list to create a breadcrumb-like effect.

:param int dir_shorten_len:
	shorten parent directory names to this length (e.g. 
	:file:`/long/path/to/powerline` → :file:`/l/p/t/powerline`)
:param int dir_limit_depth:
	limit directory depth to this number (e.g. 
	:file:`/long/path/to/powerline` → :file:`⋯/to/powerline`)
:param bool use_path_separator:
	Use path separator in place of soft divider.
:param bool shorten_home:
	Shorten home directory to ``~``.
:param str ellipsis:
	Specifies what to use in place of omitted directories. Use None to not 
	show this subsegment at all.

Divider highlight group used: ``cwd:divider``.

Highlight groups used: ``cwd:current_folder`` or ``cwd``. It is recommended to define all highlight groups.
''')


try:
	import psutil

	# psutil-2.0.0: psutil.Process.username is unbound method
	if callable(psutil.Process.username):
		def _get_user():
			return psutil.Process(os.getpid()).username()
	# pre psutil-2.0.0: psutil.Process.username has type property
	else:
		def _get_user():
			return psutil.Process(os.getpid()).username
except ImportError:
	try:
		import pwd
	except ImportError:
		from getpass import getuser as _get_user
	else:
		try:
			from os import geteuid as getuid
		except ImportError:
			from os import getuid

		def _get_user():
			return pwd.getpwuid(getuid()).pw_name


username = False
# os.geteuid is not available on windows
_geteuid = getattr(os, 'geteuid', lambda: 1)


@requires_segment_info
def user(pl, segment_info, hide_user=None, hide_domain=False):
	'''Return the current user.

	:param str hide_user:
		Omit showing segment for users with names equal to this string.
	:param bool hide_domain:
		Drop domain component if it exists in a username (delimited by '@').

	Highlights the user with the ``superuser`` if the effective user ID is 0.

	Highlight groups used: ``superuser`` or ``user``. It is recommended to define all highlight groups.
	'''
	global username
	if (
		segment_info['environ'].get('_POWERLINE_RUNNING_SHELL_TESTS')
		== 'ee5bcdc6-b749-11e7-9456-50465d597777'
	):
		return 'user'
	if username is False:
		username = _get_user()
	if username is None:
		pl.warn('Failed to get username')
		return None
	if username == hide_user:
		return None
	if hide_domain:
		try:
			username = username[:username.index('@')]
		except ValueError:
			pass
	euid = _geteuid()
	return [{
		'contents': username,
		'highlight_groups': ['user'] if euid != 0 else ['superuser', 'user'],
	}]