From 23a3dc8f3cfcd6aa0b7489a361b65191496f9100 Mon Sep 17 00:00:00 2001
From: ZyX <kp-pav@yandex.ru>
Date: Tue, 19 Aug 2014 19:05:49 +0400
Subject: [PATCH] Implement powerline reloading in libzpython bindings

---
 powerline/bindings/zsh/__init__.py   | 49 ++++++++++++++++++----------
 powerline/bindings/zsh/powerline.zsh |  2 +-
 2 files changed, 32 insertions(+), 19 deletions(-)

diff --git a/powerline/bindings/zsh/__init__.py b/powerline/bindings/zsh/__init__.py
index 1f6ae40b..ad3fa90b 100644
--- a/powerline/bindings/zsh/__init__.py
+++ b/powerline/bindings/zsh/__init__.py
@@ -1,18 +1,22 @@
 # vim:fileencoding=utf-8:noet
 from __future__ import absolute_import, unicode_literals, division, print_function
 
-import zsh
 import atexit
+
+import zsh
+
+from weakref import WeakValueDictionary, ref
+
 from powerline.shell import ShellPowerline
 from powerline.lib import parsedotval
 from powerline.lib.unicode import unicode
 
 
-used_powerlines = []
+used_powerlines = WeakValueDictionary()
 
 
 def shutdown():
-	for powerline in used_powerlines:
+	for powerline in tuple(used_powerlines.values()):
 		powerline.shutdown()
 
 
@@ -94,17 +98,28 @@ environ = Environment()
 
 
 class ZshPowerline(ShellPowerline):
+	def init(self, **kwargs):
+		super(ZshPowerline, self).init(Args(), **kwargs)
+
 	def precmd(self):
 		self.args.last_pipe_status = zsh.pipestatus()
 		self.args.last_exit_code = zsh.last_exit_code()
 
+	def do_setup(self, zsh_globals):
+		set_prompt(self, 'PS1', 'left', None, above=True)
+		set_prompt(self, 'RPS1', 'right', None)
+		set_prompt(self, 'PS2', 'left', 'continuation')
+		set_prompt(self, 'RPS2', 'right', 'continuation')
+		set_prompt(self, 'PS3', 'left', 'select')
+		used_powerlines[id(self)] = self
+		zsh_globals['_powerline'] = self
+
 
 class Prompt(object):
-	__slots__ = ('powerline', 'side', 'savedpsvar', 'savedps', 'args', 'theme', 'above')
+	__slots__ = ('powerline', 'side', 'savedpsvar', 'savedps', 'args', 'theme', 'above', '__weakref__')
 
 	def __init__(self, powerline, side, theme, savedpsvar=None, savedps=None, above=False):
 		self.powerline = powerline
-		powerline.setup(self)
 		self.side = side
 		self.above = above
 		self.savedpsvar = savedpsvar
@@ -143,9 +158,7 @@ class Prompt(object):
 	def __del__(self):
 		if self.savedps:
 			zsh.setvalue(self.savedpsvar, self.savedps)
-		used_powerlines.remove(self.powerline)
-		if self.powerline not in used_powerlines:
-			self.powerline.shutdown()
+		self.powerline.shutdown()
 
 
 def set_prompt(powerline, psvar, side, theme, above=False):
@@ -155,18 +168,18 @@ def set_prompt(powerline, psvar, side, theme, above=False):
 		savedps = None
 	zpyvar = 'ZPYTHON_POWERLINE_' + psvar
 	prompt = Prompt(powerline, side, theme, psvar, savedps, above)
+	zsh.eval('unset ' + zpyvar)
 	zsh.set_special_string(zpyvar, prompt)
 	zsh.setvalue(psvar, '${' + zpyvar + '}')
+	return ref(prompt)
 
 
-def setup():
-	powerline = ZshPowerline(Args())
-	used_powerlines.append(powerline)
-	used_powerlines.append(powerline)
-	set_prompt(powerline, 'PS1', 'left', None, above=True)
-	set_prompt(powerline, 'RPS1', 'right', None)
-	set_prompt(powerline, 'PS2', 'left', 'continuation')
-	set_prompt(powerline, 'RPS2', 'right', 'continuation')
-	set_prompt(powerline, 'PS3', 'left', 'select')
+def reload():
+	for powerline in tuple(used_powerlines.values()):
+		powerline.reload()
+
+
+def setup(zsh_globals):
+	powerline = ZshPowerline()
+	powerline.setup(zsh_globals)
 	atexit.register(shutdown)
-	return powerline
diff --git a/powerline/bindings/zsh/powerline.zsh b/powerline/bindings/zsh/powerline.zsh
index 8a7078f1..113a306d 100644
--- a/powerline/bindings/zsh/powerline.zsh
+++ b/powerline/bindings/zsh/powerline.zsh
@@ -122,7 +122,7 @@ _powerline_setup_prompt() {
 	if test -z "${POWERLINE_NO_ZSH_ZPYTHON}" && { zmodload libzpython || zmodload zsh/zpython } &>/dev/null ; then
 		precmd_functions+=( _powerline_update_counter )
 		zpython 'from powerline.bindings.zsh import setup as _powerline_setup'
-		zpython '_powerline = _powerline_setup()'
+		zpython '_powerline_setup(globals())'
 		zpython 'del _powerline_setup'
 	else
 		if test -z "${POWERLINE_COMMAND}" ; then