Add arg to optionally suppress user segment

This adds the option to common.user to suppress display if the username
matches the given string in the new "hide_user" argument.
This commit is contained in:
Kenny Root 2014-01-10 13:00:25 -08:00
parent 54e7fe91ba
commit c33d56e73c
2 changed files with 10 additions and 1 deletions

View File

@ -591,9 +591,12 @@ username = False
_geteuid = getattr(os, 'geteuid', lambda: 1)
def user(pl, segment_info=None):
def user(pl, segment_info=None, hide_user=None):
'''Return the current user.
:param str hide_user:
will suppress display of username if it matches the given string
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.
@ -604,6 +607,8 @@ def user(pl, segment_info=None):
if username is None:
pl.warn('Failed to get username')
return None
if username == hide_user:
return None
euid = _geteuid()
return [{
'contents': username,

View File

@ -80,6 +80,10 @@ class TestCommon(TestCase):
self.assertEqual(common.user(pl=pl, segment_info=segment_info), [
{'contents': 'def', 'highlight_group': 'user'}
])
self.assertEqual(common.user(pl=pl, segment_info=segment_info, hide_user='abc'), [
{'contents': 'def', 'highlight_group': 'user'}
])
self.assertEqual(common.user(pl=pl, segment_info=segment_info, hide_user='def'), None)
with replace_attr(common, '_geteuid', lambda: 0):
self.assertEqual(common.user(pl=pl, segment_info=segment_info), [
{'contents': 'def', 'highlight_group': ['superuser', 'user']}