Added support for all powerline prompts (in2, out, rewrite)

Also fixed problem with prompt not really updating.

Fixes #141

powerline/bindings/ipython/post_0_11.py

@@ -5,17 +5,26 @@ from IPython.core.prompts import PromptManager
 from IPython.core.hooks import TryNext
 
 
+class IpythonInfo(object):
+	def __init__(self, shell):
+		self._shell = shell
+
+	@property
+	def prompt_count(self):
+		return self._shell.execution_count
+
+
 class PowerlinePromptManager(PromptManager):
 	powerline = None
 
-	def __init__(self, powerline, **kwargs):
+	def __init__(self, powerline, shell):
 		self.powerline = powerline
-		super(PowerlinePromptManager, self).__init__(**kwargs)
+		self.powerline_segment_info = IpythonInfo(shell)
+		self.shell = shell
 
 	def render(self, name, color=True, *args, **kwargs):
-		if name != 'in':
+		width = None if name == 'in' else self.width
-			return super(PowerlinePromptManager, self).render(name, color, *args, **kwargs)
+		res, res_nocolor = self.powerline.renderer.render(output_raw=True, width=width, matcher_info=name, segment_info=self.powerline_segment_info)
-		res, res_nocolor = self.powerline.renderer.render(output_raw=True)
 		self.txtwidth = len(res_nocolor)
 		self.width = self.txtwidth
 		return res if color else res_nocolor
@@ -39,8 +48,7 @@ def load_ipython_extension(ip):
 	old_prompt_manager = ip.prompt_manager
 	powerline = ConfigurableIpythonPowerline(ip)
 
-	ip.prompt_manager = PowerlinePromptManager(powerline=powerline,
+	ip.prompt_manager = PowerlinePromptManager(powerline=powerline, shell=ip.prompt_manager.shell)
-		shell=ip.prompt_manager.shell, config=ip.prompt_manager.config)
 
 	def shutdown_hook():
 		powerline.renderer.shutdown()

powerline/bindings/ipython/pre_0_11.py

@@ -1,23 +1,105 @@
 # vim:fileencoding=utf-8:noet
 from powerline.ipython import IpythonPowerline
-from IPython.Prompts import BasePrompt
+from IPython.Prompts import BasePrompt, Prompt1
 from IPython.ipapi import get as get_ipython
 from IPython.ipapi import TryNext
 
+import re
+
+
+def string(s):
+	if type(s) is not str:
+		return s.encode('utf-8')
+	else:
+		return s
+
+
+# HACK: ipython tries to only leave us with plain ASCII
+class RewriteResult(object):
+	def __init__(self, prompt):
+		self.prompt = string(prompt)
+
+	def __str__(self):
+		return self.prompt
+
+	def __add__(self, s):
+		if type(s) is not str:
+			try:
+				s = s.encode('utf-8')
+			except AttributeError:
+				raise NotImplementedError
+		return RewriteResult(self.prompt + s)
+
+
+class IpythonInfo(object):
+	def __init__(self, cache):
+		self._cache = cache
+
+	@property
+	def prompt_count(self):
+		return self._cache.prompt_count
+
 
 class PowerlinePrompt(BasePrompt):
-	def __init__(self, powerline, *args, **kwargs):
+	def __init__(self, powerline, powerline_last_in, old_prompt):
 		self.powerline = powerline
-		super(PowerlinePrompt, self).__init__(*args, **kwargs)
+		self.powerline_last_in = powerline_last_in
+		self.powerline_segment_info = IpythonInfo(old_prompt.cache)
+		self.cache = old_prompt.cache
+		if hasattr(old_prompt, 'sep'):
+			self.sep = old_prompt.sep
+		self.pad_left = False
+
+	def __str__(self):
+		self.set_p_str()
+		return string(self.p_str)
+
+	def set_p_str(self, width=None):
+		self.p_str, self.p_str_nocolor = (
+			self.powerline.renderer.render(output_raw=True,
+								segment_info=self.powerline_segment_info,
+								matcher_info=self.powerline_prompt_type,
+								width=width)
+		)
+
+	@staticmethod
+	def set_colors():
+		pass
+
+
+class PowerlinePrompt1(PowerlinePrompt):
+	powerline_prompt_type = 'in'
+	rspace = re.compile(r'(\s*)$')
+
+	def __str__(self):
+		self.cache.prompt_count += 1
+		self.set_p_str()
+		self.cache.last_prompt = self.p_str_nocolor.split('\n')[-1]
+		return string(self.p_str)
 
 	def set_p_str(self):
-		self.p_str, self.p_str_nocolor = self.powerline.renderer.render(output_raw=True)
+		super(PowerlinePrompt1, self).set_p_str()
 		self.nrspaces = len(self.rspace.search(self.p_str_nocolor).group())
-		self.prompt_text_len = len(self.p_str_nocolor) - self.nrspaces - 1
+		self.prompt_text_len = len(self.p_str_nocolor) - self.nrspaces
+		self.powerline_last_in['nrspaces'] = self.nrspaces
+		self.powerline_last_in['prompt_text_len'] = self.prompt_text_len
 
 	def auto_rewrite(self):
-		# TODO color this
+		return RewriteResult(self.powerline.renderer.render(matcher_info='rewrite', width=self.prompt_text_len, segment_info=self.powerline_segment_info)
-		return '%s>%s' % ('-' * self.prompt_text_len, ' ' * self.nrspaces)
+						+ (' ' * self.nrspaces))
+
+class PowerlinePromptOut(PowerlinePrompt):
+	powerline_prompt_type = 'out'
+
+	def set_p_str(self):
+		super(PowerlinePromptOut, self).set_p_str(width=self.powerline_last_in['prompt_text_len'])
+		spaces = ' ' * self.powerline_last_in['nrspaces']
+		self.p_str += spaces
+		self.p_str_nocolor += spaces
+
+
+class PowerlinePrompt2(PowerlinePromptOut):
+	powerline_prompt_type = 'in2'
 
 
 class ConfigurableIpythonPowerline(IpythonPowerline):
@@ -28,17 +110,20 @@ class ConfigurableIpythonPowerline(IpythonPowerline):
 		super(ConfigurableIpythonPowerline, self).__init__()
 
 
-def setup(prompt='1', **kwargs):
+def setup(**kwargs):
 	ip = get_ipython()
 
 	powerline = ConfigurableIpythonPowerline(**kwargs)
 
-	attr = 'prompt' + prompt
-
 	def late_startup_hook():
-		old_prompt = getattr(ip.IP.outputcache, attr)
+		last_in = {'nrspaces': 0, 'prompt_text_len': None}
-		setattr(ip.IP.outputcache, attr, PowerlinePrompt(powerline,
+		for attr, prompt_class in (
-			old_prompt.cache, old_prompt.sep, '', old_prompt.pad_left))
+			('prompt1', PowerlinePrompt1),
+			('prompt2', PowerlinePrompt2),
+			('prompt_out', PowerlinePromptOut)
+		):
+			old_prompt = getattr(ip.IP.outputcache, attr)
+			setattr(ip.IP.outputcache, attr, prompt_class(powerline, last_in, old_prompt))
 		raise TryNext()
 
 	def shutdown_hook():

powerline/config_files/colorschemes/ipython/default.json

@@ -2,6 +2,7 @@
 	"name": "Default color scheme for IPython prompt",
 	"groups": {
 		"virtualenv": { "fg": "white", "bg": "darkcyan" },
-		"prompt": { "fg": "gray9", "bg": "gray4" }
+		"prompt": { "fg": "gray9", "bg": "gray4" },
+		"prompt_count": { "fg": "white", "bg": "gray4" }
 	}
 }

powerline/config_files/config.json

@@ -16,7 +16,12 @@
 	"ext": {
 		"ipython": {
 			"colorscheme": "default",
-			"theme": "default"
+			"theme": "in",
+			"local_themes": {
+				"rewrite": "rewrite",
+				"out": "out",
+				"in2": "in2"
+			}
 		},
 		"shell": {
 			"colorscheme": "default",

powerline/config_files/themes/ipython/default.json

@@ -1,16 +0,0 @@
-{
-	"default_module": "powerline.segments.common",
-	"segments": {
-		"left": [
-			{
-				"name": "virtualenv"
-			},
-			{
-				"type": "string",
-				"contents": "In:",
-				"highlight_group": ["prompt"],
-				"draw_divider": false
-			}
-		]
-	}
-}

powerline/config_files/themes/ipython/in.json

@@ -0,0 +1,26 @@
+{
+	"default_module": "powerline.segments.common",
+	"segments": {
+		"left": [
+			{
+				"name": "virtualenv"
+			},
+			{
+				"type": "string",
+				"contents": "In[",
+				"draw_divider": false,
+				"highlight_group": ["prompt"]
+			},
+			{
+				"name": "prompt_count",
+				"module": "powerline.segments.ipython",
+				"draw_divider": false
+			},
+			{
+				"type": "string",
+				"contents": "]",
+				"highlight_group": ["prompt"]
+			}
+		]
+	}
+}

powerline/config_files/themes/ipython/in2.json

@@ -0,0 +1,13 @@
+{
+	"default_module": "powerline.segments.common",
+	"segments": {
+		"left": [
+			{
+				"type": "string",
+				"contents": "",
+				"width": "auto",
+				"highlight_group": ["prompt"]
+			}
+		]
+	}
+}

powerline/config_files/themes/ipython/out.json

@@ -0,0 +1,25 @@
+{
+	"default_module": "powerline.segments.common",
+	"segments": {
+		"left": [
+			{
+				"type": "string",
+				"contents": "Out[",
+				"draw_divider": false,
+				"width": "auto",
+				"align": "r",
+				"highlight_group": ["prompt"]
+			},
+			{
+				"name": "prompt_count",
+				"module": "powerline.segments.ipython",
+				"draw_divider": false
+			},
+			{
+				"type": "string",
+				"contents": "]",
+				"highlight_group": ["prompt"]
+			}
+		]
+	}
+}

powerline/config_files/themes/ipython/rewrite.json

@@ -0,0 +1,23 @@
+{
+	"segments": {
+		"left": [
+			{
+				"type": "string",
+				"contents": "",
+				"draw_divider": false,
+				"width": "auto",
+				"highlight_group": ["prompt"]
+			},
+			{
+				"name": "prompt_count",
+				"module": "powerline.segments.ipython",
+				"draw_divider": false
+			},
+			{
+				"type": "string",
+				"contents": ">",
+				"highlight_group": ["prompt"]
+			}
+		]
+	}
+}

powerline/ipython.py

@@ -14,6 +14,9 @@ class IpythonPowerline(Powerline):
 		else:
 			return super(IpythonPowerline, self).get_config_paths()
 
+	def get_local_themes(self, local_themes):
+		return dict(((type, {'config': self.load_theme_config(name)}) for type, name in local_themes.items()))
+
 	def load_main_config(self):
 		r = super(IpythonPowerline, self).load_main_config()
 		if self.config_overrides:

powerline/lint/__init__.py

@@ -416,6 +416,15 @@ main_spec = (Spec(
 			local_themes=Spec()
 				.unknown_spec(lambda *args: check_matcher_func('vim', *args), theme_spec())
 		),
+		ipython=Spec(
+			colorscheme=colorscheme_spec(),
+			theme=theme_spec(),
+			local_themes=Spec(
+				in2=theme_spec(),
+				out=theme_spec(),
+				rewrite=theme_spec(),
+			),
+		),
 	).unknown_spec(check_ext,
 				Spec(
 					colorscheme=colorscheme_spec(),

powerline/renderers/ipython.py

@@ -1,6 +1,7 @@
 # vim:fileencoding=utf-8:noet
 
 from powerline.renderers.shell import ShellRenderer
+from powerline.theme import Theme
 
 
 class IpythonRenderer(ShellRenderer):
@@ -8,5 +9,16 @@ class IpythonRenderer(ShellRenderer):
 	escape_hl_start = '\x01'
 	escape_hl_end = '\x02'
 
+	def get_theme(self, matcher_info):
+		if matcher_info == 'in':
+			return self.theme
+		else:
+			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 = IpythonRenderer

powerline/renderers/vim.py

@@ -46,9 +46,11 @@ class VimRenderer(Renderer):
 		for matcher in self.local_themes.keys():
 			if matcher(matcher_info):
 				match = self.local_themes[matcher]
-				if 'config' in match:
+				try:
-					match['theme'] = Theme(theme_config=match.pop('config'), top_theme_config=self.theme_config, **self.theme_kwargs)
+					return match['theme']
-				return match['theme']
+				except KeyError:
+					match['theme'] = Theme(theme_config=match['config'], top_theme_config=self.theme_config, **self.theme_kwargs)
+					return match['theme']
 		else:
 			return self.theme
 

powerline/segments/ipython.py

@@ -0,0 +1,7 @@
+# vim:fileencoding=utf-8:noet
+
+from powerline.theme import requires_segment_info
+
+@requires_segment_info
+def prompt_count(segment_info):
+	return str(segment_info.prompt_count)

tests/test_configuration.py

@@ -101,8 +101,10 @@ class TestConfig(TestCase):
 			theme_overrides = {}
 
 		powerline = IpyPowerline()
-		powerline.renderer.render()
+		segment_info = Args(prompt_count=1)
-		powerline.renderer.render()
+		for prompt_type in ['in', 'in2', 'out', 'rewrite']:
+			powerline.renderer.render(matcher_info=prompt_type, segment_info=segment_info)
+			powerline.renderer.render(matcher_info=prompt_type, segment_info=segment_info)
 		shutdown(powerline)
 
 	def test_wm(self):