Merge pull request #809 from ZyX-I/zsh-local-themes

Add support for [R]PS2 and PS3 prompts

powerline/bindings/zsh/__init__.py

@@ -55,8 +55,8 @@ class Args(object):
 
 	@property
 	def jobnum(self):
-		zsh.eval('integer POWERLINE_JOBNUM=${(%):-%j}')
-		return zsh.getvalue('POWERLINE_JOBNUM')
+		zsh.eval('integer _POWERLINE_JOBNUM=${(%):-%j}')
+		return zsh.getvalue('_POWERLINE_JOBNUM')
 
 
 def string(s):
@@ -94,20 +94,29 @@ environ = Environment()
 
 
 class Prompt(object):
-	__slots__ = ('powerline', 'side', 'savedpsvar', 'savedps', 'args')
+	__slots__ = ('powerline', 'side', 'savedpsvar', 'savedps', 'args', 'theme')
 
-	def __init__(self, powerline, side, savedpsvar=None, savedps=None):
+	def __init__(self, powerline, side, theme, savedpsvar=None, savedps=None):
 		self.powerline = powerline
 		self.side = side
 		self.savedpsvar = savedpsvar
 		self.savedps = savedps
 		self.args = powerline.args
+		self.theme = theme
 
 	def __str__(self):
+		zsh.eval('_POWERLINE_PARSER_STATE="${(%):-%_}"')
+		segment_info = {
+			'args': self.args,
+			'environ': environ,
+			'client_id': 1,
+			'local_theme': self.theme,
+			'parser_state': zsh.getvalue('_POWERLINE_PARSER_STATE'),
+		}
 		r = self.powerline.render(
 			width=zsh.columns(),
 			side=self.side,
-			segment_info={'args': self.args, 'environ': environ}
+			segment_info=segment_info,
 		)
 		if type(r) is not str:
 			if type(r) is bytes:
@@ -124,10 +133,13 @@ class Prompt(object):
 			self.powerline.shutdown()
 
 
-def set_prompt(powerline, psvar, side):
+def set_prompt(powerline, psvar, side, theme):
+	try:
 		savedps = zsh.getvalue(psvar)
+	except IndexError:
+		savedps = None
 	zpyvar = 'ZPYTHON_POWERLINE_' + psvar
-	prompt = Prompt(powerline, side, psvar, savedps)
+	prompt = Prompt(powerline, side, theme, psvar, savedps)
 	zsh.set_special_string(zpyvar, prompt)
 	zsh.setvalue(psvar, '${' + zpyvar + '}')
 
@@ -136,6 +148,9 @@ def setup():
 	powerline = ShellPowerline(Args())
 	used_powerlines.append(powerline)
 	used_powerlines.append(powerline)
-	set_prompt(powerline, 'PS1', 'left')
-	set_prompt(powerline, 'RPS1', 'right')
+	set_prompt(powerline, 'PS1', 'left', None)
+	set_prompt(powerline, 'RPS1', 'right', None)
+	set_prompt(powerline, 'PS2', 'left', 'continuation')
+	set_prompt(powerline, 'RPS2', 'right', 'continuation')
+	set_prompt(powerline, 'PS3', 'left', 'select')
 	atexit.register(shutdown)

powerline/bindings/zsh/powerline.zsh

@@ -113,9 +113,14 @@ _powerline_setup_prompt() {
 		zpython 'del _powerline_setup'
 	else
 		local add_args='--last_exit_code=$? --last_pipe_status="$pipestatus"'
+		add_args+=' --renderer_arg="client_id=$$"'
 		add_args+=' --jobnum=$_POWERLINE_JOBNUM'
+		local add_args_2=$add_args' -R parser_state=${(%%):-%_} -R local_theme=continuation'
 		PS1='$($POWERLINE_COMMAND shell left -r zsh_prompt '$add_args')'
 		RPS1='$($POWERLINE_COMMAND shell right -r zsh_prompt '$add_args')'
+		PS2='$($POWERLINE_COMMAND shell left -r zsh_prompt '$add_args_2')'
+		RPS2='$($POWERLINE_COMMAND shell right -r zsh_prompt '$add_args_2')'
+		PS3='$($POWERLINE_COMMAND shell left -r zsh_prompt -R local_theme=select '$add_args')'
 	fi
 }
 

powerline/config_files/colorschemes/shell/default.json

@@ -8,6 +8,8 @@
 		"branch": { "fg": "gray9", "bg": "gray2" },
 		"branch_dirty": { "fg": "brightyellow", "bg": "gray2" },
 		"branch_clean": { "fg": "gray9", "bg": "gray2" },
+		"continuation": { "fg": "gray9", "bg": "gray4" },
+		"continuation:current": { "fg": "gray10", "bg": "gray4", "attr": ["bold"] },
 		"cwd": { "fg": "gray9", "bg": "gray4" },
 		"cwd:current_folder": { "fg": "gray10", "bg": "gray4", "attr": ["bold"] },
 		"cwd:divider": { "fg": "gray7", "bg": "gray4" },

powerline/config_files/colorschemes/shell/solarized.json

@@ -8,6 +8,8 @@
 		"branch":               { "fg": "gray61", "bg": "royalblue5" },
 		"branch_dirty":         { "fg": "yellow", "bg": "royalblue5" },
 		"branch_clean":         { "fg": "gray61", "bg": "royalblue5" },
+		"continuation":         { "fg": "lightyellow", "bg": "darkgreencopper" },
+		"continuation:current": { "fg": "oldlace", "bg": "darkgreencopper", "attr": ["bold"] },
 		"cwd":                  { "fg": "lightyellow", "bg": "darkgreencopper" },
 		"cwd:current_folder":   { "fg": "oldlace", "bg": "darkgreencopper", "attr": ["bold"] },
 		"cwd:divider":          { "fg": "gray61", "bg": "darkgreencopper" },

powerline/config_files/config.json

@@ -25,7 +25,11 @@
 		},
 		"shell": {
 			"colorscheme": "default",
-			"theme": "default"
+			"theme": "default",
+			"local_themes": {
+				"continuation": "continuation",
+				"select": "select"
+			}
 		},
 		"tmux": {
 			"colorscheme": "default",

powerline/config_files/themes/shell/continuation.json

@@ -0,0 +1,12 @@
+{
+	"default_module": "powerline.segments.shell",
+	"segments": {
+		"left": [
+			{
+				"name": "continuation"
+			}
+		],
+		"right": [
+		]
+	}
+}

powerline/config_files/themes/shell/select.json

@@ -0,0 +1,13 @@
+{
+	"segments": {
+		"left": [
+			{
+				"type": "string",
+				"contents": "Select variant",
+				"width": "auto",
+				"align": "r",
+				"highlight_group": ["continuation:current"]
+			}
+		]
+	}
+}

powerline/lint/__init__.py

@@ -478,6 +478,14 @@ main_spec = (Spec(
 				rewrite=theme_spec(),
 			),
 		).optional(),
+		shell=Spec(
+			colorscheme=colorscheme_spec(),
+			theme=theme_spec(),
+			local_themes=Spec(
+				continuation=theme_spec(),
+				select=theme_spec(),
+			),
+		).optional(),
 	).unknown_spec(check_ext,
 				Spec(
 					colorscheme=colorscheme_spec(),

powerline/renderers/ipython.py

@@ -22,7 +22,11 @@ class IpythonRenderer(ShellRenderer):
 			try:
 				return match['theme']
 			except KeyError:
-				match['theme'] = Theme(theme_config=match['config'], top_theme_config=self.theme_config, **self.theme_kwargs)
+				match['theme'] = Theme(
+					theme_config=match['config'],
+					top_theme_config=self.theme_config,
+					**self.theme_kwargs
+				)
 				return match['theme']
 
 	def shutdown(self):

powerline/renderers/zsh_prompt.py

@@ -3,6 +3,7 @@
 from __future__ import absolute_import, unicode_literals
 
 from powerline.renderers.shell import ShellRenderer
+from powerline.theme import Theme
 
 
 class ZshPromptRenderer(ShellRenderer):
@@ -13,5 +14,47 @@ class ZshPromptRenderer(ShellRenderer):
 	character_translations = ShellRenderer.character_translations.copy()
 	character_translations[ord('%')] = '%%'
 
+	old_widths = {}
+
+	def render(self, segment_info, *args, **kwargs):
+		client_id = segment_info.get('client_id')
+		key = (client_id, kwargs.get('side'))
+		kwargs = kwargs.copy()
+		width = kwargs.pop('width', None)
+		local_theme = segment_info.get('local_theme')
+		if client_id and local_theme:
+			output_raw = False
+			try:
+				width = self.old_widths[key]
+			except KeyError:
+				pass
+		else:
+			output_raw = True
+		ret = super(ShellRenderer, self).render(
+			output_raw=output_raw,
+			width=width,
+			matcher_info=local_theme,
+			segment_info=segment_info,
+			*args, **kwargs
+		)
+		if output_raw:
+			self.old_widths[key] = len(ret[1])
+			ret = ret[0]
+		return ret
+
+	def get_theme(self, matcher_info):
+		if not matcher_info:
+			return self.theme
+		match = self.local_themes[matcher_info]
+		try:
+			return match['theme']
+		except KeyError:
+			match['theme'] = Theme(
+				theme_config=match['config'],
+				top_theme_config=self.theme_config,
+				**self.theme_kwargs
+			)
+			return match['theme']
+
 
 renderer = ZshPromptRenderer

powerline/segments/shell.py

@@ -72,3 +72,47 @@ def mode(pl, segment_info, override={'vicmd': 'COMMND', 'viins': 'INSERT'}, defa
 		# code or by somebody knowing what he is doing there is absolutely no 
 		# need in keeping translations dictionary.
 		return mode.upper()
+
+
+@requires_segment_info
+def continuation(pl, segment_info, omit_cmdsubst=True, right_align=False, renames={}):
+	'''Display parser state.
+
+	:param bool omit_cmdsubst:
+		Do not display cmdsubst parser state if it is the last one.
+	:param bool right_align:
+		Align to the right.
+	:param dict renames:
+		Rename states: ``{old_name : new_name}``. If ``new_name`` is ``None`` 
+		then given state is not displayed.
+
+	Highlight groups used: ``continuation``, ``continuation:current``.
+	'''
+	if not segment_info.get('parser_state'):
+		return None
+	ret = []
+
+	for state in segment_info['parser_state'].split():
+		state = renames.get(state, state)
+		if state:
+			ret.append({
+				'contents': state,
+				'highlight_group': 'continuation',
+				'draw_inner_divider': True,
+			})
+
+	if omit_cmdsubst and ret[-1]['contents'] == 'cmdsubst':
+		ret.pop(-1)
+
+	if not ret:
+		ret.append({
+			'contents': ''
+		})
+
+	if right_align:
+		ret[0].update(width='auto', align='r')
+		ret[-1]['highlight_group'] = 'continuation:current'
+	else:
+		ret[-1].update(width='auto', align='l', highlight_group='continuation:current')
+
+	return ret

powerline/shell.py

@@ -38,6 +38,13 @@ class ShellPowerline(Powerline):
 		else:
 			return super(ShellPowerline, self).get_config_paths()
 
+	def get_local_themes(self, local_themes):
+		if not local_themes:
+			return {}
+
+		return dict(((key, {'config': self.load_theme_config(val)})
+					for key, val in local_themes.items()))
+
 
 def get_argparser(parser=None, *args, **kwargs):
 	if not parser:
@@ -54,7 +61,7 @@ def get_argparser(parser=None, *args, **kwargs):
 	p.add_argument('-c', '--config', metavar='KEY.KEY=VALUE', action='append')
 	p.add_argument('-t', '--theme_option', metavar='THEME.KEY.KEY=VALUE', action='append')
 	p.add_argument('-p', '--config_path', metavar='PATH')
-	p.add_argument('-R', '--renderer_arg', metavar='KEY="VAL"', type=lambda a: dict([parsedotval(a)]))
+	p.add_argument('-R', '--renderer_arg', metavar='KEY=VAL', action='append')
 	return p
 
 
@@ -65,3 +72,5 @@ def finish_args(args):
 		args.theme_option = mergeargs((parsedotval(v) for v in args.theme_option))
 	else:
 		args.theme_option = {}
+	if args.renderer_arg:
+		args.renderer_arg = mergeargs((parsedotval(v) for v in args.renderer_arg))

tests/test_cmdline.py

@@ -43,7 +43,6 @@ class TestParser(TestCase):
 				(['-r', 'zsh_prompt'],                   'too few arguments|the following arguments are required: ext'),
 				(['shell', '--last_exit_code', 'i'],     'invalid int value'),
 				(['shell', '--last_pipe_status', '1 i'], 'invalid <lambda> value'),
-				(['shell', '-R', 'abc'],                 'invalid <lambda> value'),
 			]:
 				self.assertRaises(SystemExit, parser.parse_args, raising_args)
 				self.assertFalse(out.getvalue())

tests/test_configuration.py

@@ -85,10 +85,18 @@ class TestConfig(TestCase):
 	def test_zsh(self):
 		from powerline.shell import ShellPowerline
 		args = Args(last_pipe_status=[1, 0], jobnum=0, ext=['shell'], renderer_module='zsh_prompt')
+		segment_info = {'args': args}
 		with ShellPowerline(args, run_once=False) as powerline:
-			powerline.render(segment_info={'args': args})
+			powerline.render(segment_info=segment_info)
 		with ShellPowerline(args, run_once=False) as powerline:
-			powerline.render(segment_info={'args': args})
+			powerline.render(segment_info=segment_info)
+		segment_info['local_theme'] = 'select'
+		with ShellPowerline(args, run_once=False) as powerline:
+			powerline.render(segment_info=segment_info)
+		segment_info['local_theme'] = 'continuation'
+		segment_info['parser_state'] = 'if cmdsubst'
+		with ShellPowerline(args, run_once=False) as powerline:
+			powerline.render(segment_info=segment_info)
 
 	def test_bash(self):
 		from powerline.shell import ShellPowerline

tests/test_segments.py

@@ -48,6 +48,103 @@ class TestShell(TestCase):
 		self.assertEqual(shell.jobnum(pl=pl, segment_info=segment_info, show_zero=False), '1')
 		self.assertEqual(shell.jobnum(pl=pl, segment_info=segment_info, show_zero=True), '1')
 
+	def test_continuation(self):
+		pl = Pl()
+		self.assertEqual(shell.continuation(pl=pl, segment_info={}), None)
+		segment_info = {'parser_state': 'if cmdsubst'}
+		self.assertEqual(shell.continuation(pl=pl, segment_info=segment_info), [
+			{
+				'contents': 'if',
+				'draw_inner_divider': True,
+				'highlight_group': 'continuation:current',
+				'width': 'auto',
+				'align': 'l',
+			},
+		])
+		self.assertEqual(shell.continuation(pl=pl, segment_info=segment_info, right_align=True), [
+			{
+				'contents': 'if',
+				'draw_inner_divider': True,
+				'highlight_group': 'continuation:current',
+				'width': 'auto',
+				'align': 'r',
+			},
+		])
+		self.assertEqual(shell.continuation(pl=pl, segment_info=segment_info, omit_cmdsubst=False), [
+			{
+				'contents': 'if',
+				'draw_inner_divider': True,
+				'highlight_group': 'continuation',
+			},
+			{
+				'contents': 'cmdsubst',
+				'draw_inner_divider': True,
+				'highlight_group': 'continuation:current',
+				'width': 'auto',
+				'align': 'l',
+			},
+		])
+		self.assertEqual(shell.continuation(pl=pl, segment_info=segment_info, omit_cmdsubst=False, right_align=True), [
+			{
+				'contents': 'if',
+				'draw_inner_divider': True,
+				'highlight_group': 'continuation',
+				'width': 'auto',
+				'align': 'r',
+			},
+			{
+				'contents': 'cmdsubst',
+				'draw_inner_divider': True,
+				'highlight_group': 'continuation:current',
+			},
+		])
+		self.assertEqual(shell.continuation(pl=pl, segment_info=segment_info, omit_cmdsubst=True, right_align=True), [
+			{
+				'contents': 'if',
+				'draw_inner_divider': True,
+				'highlight_group': 'continuation:current',
+				'width': 'auto',
+				'align': 'r',
+			},
+		])
+		self.assertEqual(shell.continuation(pl=pl, segment_info=segment_info, omit_cmdsubst=True, right_align=True, renames={'if': 'IF'}), [
+			{
+				'contents': 'IF',
+				'draw_inner_divider': True,
+				'highlight_group': 'continuation:current',
+				'width': 'auto',
+				'align': 'r',
+			},
+		])
+		self.assertEqual(shell.continuation(pl=pl, segment_info=segment_info, omit_cmdsubst=True, right_align=True, renames={'if': None}), [
+			{
+				'contents': '',
+				'highlight_group': 'continuation:current',
+				'width': 'auto',
+				'align': 'r',
+			},
+		])
+		segment_info = {'parser_state': 'then then then cmdsubst'}
+		self.assertEqual(shell.continuation(pl=pl, segment_info=segment_info), [
+			{
+				'contents': 'then',
+				'draw_inner_divider': True,
+				'highlight_group': 'continuation',
+			},
+			{
+				'contents': 'then',
+				'draw_inner_divider': True,
+				'highlight_group': 'continuation',
+			},
+			{
+				'contents': 'then',
+				'draw_inner_divider': True,
+				'highlight_group': 'continuation:current',
+				'width': 'auto',
+				'align': 'l',
+			},
+		])
+
 
 class TestCommon(TestCase):
 	def test_hostname(self):

tests/test_shells/input.zsh

@@ -10,6 +10,12 @@ VIRTUAL_ENV="$HOME/.virtenvs/some-virtual-environment"
 VIRTUAL_ENV=
 bash -c 'echo $$>pid ; while true ; do sleep 0.1s ; done' &
 false
+select abc in def ghi jkl
+do
+	echo $abc
+	break
+done
+1
 kill `cat pid` ; sleep 1s
 cd "$DIR1"
 cd ../"$DIR2"

tests/test_shells/zsh.ok

@@ -6,7 +6,15 @@
   HOSTNAME  USER   BRANCH  ⋯  tests  shell  3rd  bash -c 'echo $$>pid ; while true ; do sleep 0.1s ; done' &
 [1] PID
   HOSTNAME  USER   BRANCH  ⋯  tests  shell  3rd  1  false
-  HOSTNAME  USER   BRANCH  ⋯  tests  shell  3rd  1  1  kill `cat pid` ; sleep 1s
+  HOSTNAME  USER   BRANCH  ⋯  tests  shell  3rd  1  1  select abc in def ghi jkl
+ select  do
+ select        echo $abc
+ select        break
+ select  done
+1) def  2) ghi  3) jkl  
+ Select variant  1
+def
+  HOSTNAME  USER   BRANCH  ⋯  tests  shell  3rd  1  kill `cat pid` ; sleep 1s
 [1]  + terminated  bash -c 'echo $$>pid ; while true ; do sleep 0.1s ; done'
   HOSTNAME  USER   BRANCH  ⋯  tests  shell  3rd  cd "$DIR1"
   HOSTNAME  USER   BRANCH  ⋯  shell  3rd  ^[[32m  cd ../"$DIR2"