From cd3dd70b45babcfe1a13643f2e31d8aa90b28f79 Mon Sep 17 00:00:00 2001
From: Martinez <overwatch@covalent.cz>
Date: Sat, 26 Dec 2015 01:48:43 +0100
Subject: [PATCH] new render program prototype mostly finished

---
 render.py | 571 +++++++++++++++++++++++++++++++++++++++++-------------
 zsh-init  |  59 +-----
 2 files changed, 441 insertions(+), 189 deletions(-)

diff --git a/render.py b/render.py
index bd395bc..1215389 100755
--- a/render.py
+++ b/render.py
@@ -1,213 +1,518 @@
 #! /usr/bin/env python3
 # encoding: utf-8
 
+import datetime
 import os
 import socket
+import subprocess
 import sys
 
-SYM_GIT = "\ue0a0"
-SYM_LOCK = "\ue0a2"
+COLOR_LOAD_IDLE = 10
+COLOR_LOAD_OK = 2
+COLOR_LOAD_WARN = 3
+COLOR_LOAD_ERROR = 1
+COLOR_USER_USER = 10
+COLOR_USER_ROOT = 196
+COLOR_SESSION_LOCAL = 7
+COLOR_SESSION_REMOTE = 202
+COLOR_TERM_SCREEN = 27
 
-def style_color(fg, bg):
-	return '\033[38;5;%dm\033[48;5;%dm' %(fg, bg)
+COLOR_CLOCK = 27
+COLOR_MEMSWAP = 8
+COLOR_SPACE_OK = 2
+COLOR_SPACE_WARN = 3
+COLOR_SPACE_ERROR = 1
+
+COLOR_DIR_RW = 34
+COLOR_DIR_RO = 202
+COLOR_DIR_NO = 196
+
+COLOR_GIT_CLEAN = 10
+COLOR_GIT_DIRTY = 1
+COLOR_GIT_MERGE = 196
+COLOR_GIT_UNTRACKED = 202
+COLOR_GIT_MODIFIED = 1
+COLOR_GIT_STAGED = 10
+
+COLOR_PROMPT_OK = 14
+COLOR_PROMPT_ERROR = 1
+
+MOUNT_IGNORE_FS = ["iso9660", "tmpfs"]
+MOUNT_IGNORE_DIR = ["/dev", "/proc", "/sys"]
+
+GIT_BLACKLIST = ["/var/paludis/repositories"]
+
+def style_color(fg, bg=0):
+	return "\033[38;5;%dm\033[48;5;%dm" %(fg, bg)
 
 def style_bold():
-	return '\033[1m'
+	return "\033[1m"
 
 def style_reset():
-	return '\033[0m'
+	return "\033[0m"
 
-class LoginPart:
-	user_fg = 34
-	root_fg = 196
-	remote_fg = 202
-	local_fg = 7
-	screen_fg = 27
-	bg = 0
+def colored_strlen(raw):
+	"""
+	Returns string length without ANSI control escapes.
+	"""
 	
+	l = 0
+	reading_escape = False
+	
+	for c in raw:
+		if c == "\x1b":
+			reading_escape = True
+			continue
+		if reading_escape and c == "m":
+			reading_escape = False
+			continue
+		
+		if not reading_escape:
+			l += 1
+	
+	return l
+
+def si_number(raw):
+	if raw > 2**40:
+		raw /= 2**40
+		si = "Ti"
+	elif raw > 2**30:
+		raw /= 2**30
+		si = "Gi"
+	elif raw > 2**20:
+		raw /= 2**20
+		si = "Mi"
+	elif raw > 2**10:
+		raw /= 2**10
+		si = "ki"
+	else:
+		si = ""
+	
+	return (raw, si)
+
+def space_string(free):
+	"""
+	Returns a formatting string suitable for printing out free space: float for >1k, int otherwise.
+	"""
+	
+	return "%.2f%s" if free > 1024 else "%d%s"
+
+class Sysload:
 	def __init__(self):
-		self.user = os.getlogin()
-		self.remote = bool(os.getenv('SSH_CLIENT'))
-		self.sign = ' ⇄ ' if self.remote else '@'
-		self.host = socket.gethostname()
-		self.screen = os.getenv('WINDOW')
+		loadavg = open("/proc/loadavg").read().split()
+		cores = os.sysconf(os.sysconf_names["SC_NPROCESSORS_ONLN"])
+		
+		self.load1 = float(loadavg[0]) / cores
+		self.load5 = float(loadavg[1]) / cores
+		self.load15 = float(loadavg[2]) / cores
+		self.tasks_ready, self.tasks_total = [int(x) for x in loadavg[3].split("/")]
+
+
+class Settings:
+	def __init__(self, argv):
+		try:
+			self.return_value = int(argv[1])
+			self.term_width = int(argv[2])
+			self.load_warn = float(argv[3])
+			self.load_error = float(argv[4])
+			self.space_warn = float(argv[5])
+			self.space_error = float(argv[6])
+			self.is_root = os.geteuid() == 0
+		except:
+			print("Usage: %s return_value $COLUMNS load_warn load_error space_warn space_error")
+			print("       e.g. %s 1 $COLUMNS 1.25 2.5 0.15 0.05")
+			print()
+			
+			raise
 	
-	@property
-	def length(self):
-		return len(self.user) + len(self.host) + len(self.screen) + len(self.sign) + 1
+	def get_load_color(self, load):
+		if load > self.load_error:
+			return COLOR_LOAD_ERROR
+		elif load > self.load_warn:
+			return COLOR_LOAD_WARN
+		elif load > 0.1:
+			return COLOR_LOAD_OK
+		else:
+			return COLOR_LOAD_IDLE
+	
+	def get_space_color(self, free, total):
+		free = free / total
+		
+		if free < self.space_error:
+			return COLOR_SPACE_ERROR
+		elif free < self.space_warn:
+			return COLOR_SPACE_WARN
+		else:
+			return COLOR_SPACE_OK
+
+class Part:
+	def __init__(self):
+		self.fragments = []
 	
 	def __str__(self):
-		output = ''
+		return "".join(self.fragments)
+	
+	def __len__(self):
+		return colored_strlen(str(self))
+
+class LoginPart(Part):
+	"""
+	username@hostname:screen
+	- username is green for users, red for root
+	- @ is a literal "@" if this is a local session, " ⇄ " when remote (ssh)
+	- hostname is colored according to sysload
+	- screen is the screen window ID if available
+	"""
+	
+	def __init__(self, settings, sysload):
+		Part.__init__(self)
+		is_remote = bool(os.getenv("SSH_CLIENT"))
+		screen = os.getenv("WINDOW")
 		
-		if os.geteuid() == 0:
-			output += style_color(self.root_fg, self.bg)
+		# username
+		if settings.is_root:
+			self.fragments.append(style_color(COLOR_USER_ROOT))
 		else:
-			output += style_color(self.user_fg, self.bg)
+			self.fragments.append(style_color(COLOR_USER_USER))
 		
-		output += self.user
+		self.fragments.append(style_bold())
+		user = os.getlogin()
+		self.fragments.append(user)
+		self.fragments.append(style_reset())
 		
-		if self.remote:
-			output += style_color(self.remote_fg, self.bg)
+		# sign
+		if is_remote:
+			self.fragments.append(style_color(COLOR_SESSION_REMOTE))
 		else:
-			output += style_color(self.local_fg, self.bg)
+			self.fragments.append(style_color(COLOR_SESSION_LOCAL))
 		
-		output += self.sign
-		output += self.host
+		sign = " ⇄ " if is_remote else "@"
+		self.fragments.append(sign)
 		
-		if self.screen:
-			output += style_reset()
-			output += ':'
-			output += style_color(self.screen_fg, self.bg)
-			output += self.screen
+		# hostname
+		self.fragments.append(style_color(settings.get_load_color(sysload.load1)))
+		hostname = socket.gethostname()
+		self.fragments.append(hostname)
 		
-		output += style_reset()
-		
-		return output
+		# screen window ID
+		if screen:
+			self.fragments.append(style_reset())
+			self.fragments.append(":")
+			self.fragments.append(style_bold())
+			self.fragments.append(style_color(COLOR_TERM_SCREEN))
+			self.fragments.append(screen)
+			self.fragments.append(style_reset())
 
 def list_homes():
-	with open('/etc/passwd') as f:
+	with open("/etc/passwd") as f:
 		for line in f:
-			tokens = line.split(':')
+			tokens = line.split(":")
 			
-			if tokens[5] != '/dev/null':
+			if tokens[5] != "/dev/null":
 				yield (tokens[5], tokens[0])
 
 class Dir:
-	writable_fg = 34
-	fg = 196
-	bg = 0
-	
-	def __init__(self, path, name=None, slash='/'):
+	def __init__(self, path, name=None, slash="/"):
 		self.path = path
 		self.slash = slash
-		self.lock = not os.access(path, os.W_OK)
 		self.truncated = False
+		self.name = os.path.basename(path) if name is None else name
 		
-		if name is None:
-			name = os.path.basename(path)
-		
-		self.raw_str = name + ('' if self.lock else '')
+		if os.access(path, os.W_OK):
+			self.color = COLOR_DIR_RW
+		elif os.access(path, os.X_OK):
+			self.color = COLOR_DIR_RO
+		else:
+			self.color = COLOR_DIR_NO
 	
-	@property
-	def length(self):
-		return len(self.raw_str) + len(self.slash)
+	def __len__(self):
+		return len(self.name) + len(self.slash)
 	
 	def truncate(self):
 		self.truncated = True
 		
-		if self.raw_str:
-			self.raw_str = self.raw_str[0]
+		if self.name:
+			self.name = self.name[0] + "…"
 	
 	def __str__(self):
-		output = self.slash
+		output = [self.slash]
 		
 		if self.truncated:
-			output += style_bold()
+			output.append(style_bold())
 		
-		if self.raw_str:
-			output += style_color(self.fg if self.lock else self.writable_fg, self.bg)
-			output += self.raw_str
+		output.append(style_color(self.color))
+		output.append(self.name)
 		
-		output += style_reset()
+		output.append(style_reset())
 		
-		return output
+		return "".join(output)
 
-class PathPart:
+class PathPart(Part):
+	"""
+	/path/to/cwd
+	
+	Contructs a list of Dirs from the root to CWD. When fit is called, some Dirs may be shortened.
+	"""
 	def __init__(self, term_width):
+		Part.__init__(self)
 		self.term_width = term_width
-		self.segments = []
+		self.dirs = []
 		
 		homes = dict(list_homes())
-		dirs = os.getcwd().split('/')[1:]
-		path = ''
+		dirs = os.getcwd().split("/")[1:]
+		path = ""
 		
 		for dir in dirs:
-			path += '/' + dir
+			path += "/" + dir
 			
 			if path in homes:
-				self.segments = []
+				self.dirs = []
 				
 				if homes[path] == os.getlogin():
-					self.segments.append(Dir(path, '', '~'))
+					self.dirs.append(Dir(path, "", "~"))
 				else:
-					self.segments.append(Dir(path, homes[path], '~'))
+					self.dirs.append(Dir(path, homes[path], "~"))
 			
 			else:
-				self.segments.append(Dir(path))
+				self.dirs.append(Dir(path))
 	
 	@property
-	def length(self):
-		return sum(s.length for s in self.segments)
+	def full_length(self):
+		return sum(len(dir) for dir in self.dirs)
 	
-	def truncate(self, length):
-		for s in self.segments:
-			if self.length > length:
-				s.truncate()
-			else:
-				break
-	
-	def truncate_fit(self, line):
-		space = self.term_width - sum(part.length for part in line if part is not self)
-		self.truncate(space)
-	
-	def __str__(self):
-		return ''.join(str(s) for s in self.segments)
+	def shrink_fit(self, line):
+		self.fragments = []
+		
+		available_space = self.term_width - sum(len(part) for part in line if part is not self) - len(line) + 1 # account for single-space separators
+		
+		for dir in self.dirs:
+			if self.full_length > available_space:
+				dir.truncate()
+			
+			self.fragments.append(str(dir))
 
-class GitPart:
+def command(cmd):
+	"""
+	Executes a command, returns stdout, suppresses stderr."
+	"""
+	
+	s = subprocess.Popen(cmd.split(), stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+	s.wait()
+	
+	return s.stdout.read().decode("utf-8")
+
+class GitPart(Part):
+	"""
+	↘2 ↗4 ⚠master ?11 ✎6 ✉10
+	
+	- 2 commits are available for pulling (remote is ahead)
+	- 4 commits are available for pushing (local is ahead)
+	- we're merging
+	- branch name is master
+	- 11 untracked filed
+	- 6 modified files
+	- 10 modified and staged files
+	
+	Numeric elements are hidden if the value is 0.
+	"""
+	
 	def __init__(self):
-		...
-	
-	@property
-	def length(self):
-		return len(str(self))
-	
-	def __str__(self):
-		return 'GIT'
+		Part.__init__(self)
+		
+		branch_name = command("git name-rev --name-only --no-undefined --always HEAD").strip()
+		
+		if branch_name:
+			self.fragments.append("| ")
+			
+			count_ahead = command("git log --oneline @{u}..").count("\n")
+			count_behind = command("git log --oneline ..@{u}").count("\n")
+			git_dir = command("git rev-parse --git-dir").strip()
+			merging = os.path.exists(os.path.join(git_dir, "MERGE_HEAD"))
+			untracked = command("git ls-files --other --exclude-standard").count("\n")
+			modified = command("git diff --name-only").count("\n")
+			staged = command("git diff --name-only --staged").count("\n")
+			
+			if count_ahead:
+				self.fragments.append("↘%d " %(count_ahead))
+			
+			if count_behind:
+				self.fragments.append("↗%d " %(count_behind))
+			
+			if merging:
+				self.fragments.append(style_color(COLOR_GIT_MERGE))
+				self.fragments.append(style_bold())
+				self.fragments.append("⚠")
+			elif modified or staged:
+				self.fragments.append(style_color(COLOR_GIT_DIRTY))
+			else:
+				self.fragments.append(style_color(COLOR_GIT_CLEAN))
+			
+			self.fragments.append(branch_name)
+			self.fragments.append(style_reset())
+			
+			if untracked:
+				self.fragments.append(style_color(COLOR_GIT_UNTRACKED))
+				self.fragments.append(" ?%d" %(untracked))
+			
+			if modified:
+				self.fragments.append(style_color(COLOR_GIT_MODIFIED))
+				self.fragments.append(" ✎%d" %(modified))
+			
+			if staged:
+				self.fragments.append(style_color(COLOR_GIT_STAGED))
+				self.fragments.append(" ✉%d" %(staged))
+			
+			self.fragments.append(style_reset())
 
-class Padding:
+class Padding(Part):
 	def __init__(self, term_width):
+		Part.__init__(self)
 		self.term_width = term_width
-		self.length = 0
 	
 	def expand_fit(self, line):
-		self.length = self.term_width - sum(part.length for part in line if part is not self) - len(line) + 1
-	
-	def __str__(self):
-		return ' ' * self.length
+		length = self.term_width - sum(len(part) for part in line if part is not self) - len(line) + 1 # account for single-space separators
+		
+		self.fragments = [" " * length]
 
-class StatsPart:
-	clock = 27
-	space_good = 34
-	space_warn = 208
-	space_bad = 196
-	
-	def __init__(self, red_thresh, yellow_thresh, warning_only=True):
-		self.red = red_thresh
-		self.yellow = yellow_thresh
-	
-	@property
-	def length(self):
-		return len(str(self))
-	
-	def __str__(self):
-		return '[ 13:34 | 0.53 | r4.20M'
+class StatsPart(Part):
+	"""
+	[ clock | sysload tasks | m:memory (?s:swap?) mountpoint_spaces...
+	- clock shows current day and time (in DD HH:MM)
+	- sysload is current loadavg divided by amount of cores (in 1)
+	- tasks is the total amount of tasks on the system
+	- memory shows current free memory (in octets)
+	- swap shows remaining swap space but only if any is actually used
+	- mountpoint_spaces list one free space per mountpoint
+	"""
+	def __init__(self, settings, sysload):
+		Part.__init__(self)
+		
+		# clock
+		self.fragments.append(datetime.datetime.now().strftime("[ %%d \033[38;5;%dm%%H:%%M\033[0m | " %(COLOR_CLOCK)))
+		
+		# sysload
+		self.fragments.append(style_color(settings.get_load_color(sysload.load1)))
+		self.fragments.append("%.2f " %(sysload.load1))
+		self.fragments.append(style_reset())
+		
+		# tasks
+		self.fragments.append("%d | " %(sysload.tasks_total))
+		
+		# memory (and swap, if used)
+		mem_total = 0
+		mem_free = 0
+		swap_total = 0
+		swap_free = 0
+		
+		with open("/proc/meminfo") as f:
+			for line in f:
+				if line.startswith("MemTotal"):
+					mem_total = int(line.split()[1]) * 1024
+				elif line.startswith("MemAvailable"):
+					mem_free = int(line.split()[1]) * 1024
+				elif line.startswith("SwapTotal"):
+					swap_total = int(line.split()[1]) * 1024
+				elif line.startswith("SwapFree"):
+					swap_free = int(line.split()[1]) * 1024
+		
+		self.fragments.append(style_color(COLOR_MEMSWAP))
+		self.fragments.append("m")
+		self.fragments.append(style_color(settings.get_space_color(mem_free, mem_total)))
+		mem_free_si = si_number(mem_free)
+		self.fragments.append(space_string(mem_free) %(mem_free_si[0], mem_free_si[1]))
+		self.fragments.append(style_reset())
+		
+		if swap_total - swap_free > 2**20:
+			self.fragments.append(style_color(COLOR_MEMSWAP))
+			self.fragments.append(" s")
+			self.fragments.append(style_color(settings.get_space_color(swap_free, swap_total)))
+			swap_free_si = si_number(swap_free)
+			self.fragments.append(space_string(swap_free) %(swap_free_si[0], swap_free_si[1]))
+			self.fragments.append(style_reset())
+		
+		# mountpoints
+		names = []
+		
+		with open("/proc/self/mounts") as f:
+			for line in f:
+				_, dir, type, options, *rest = line.split()
+				
+				# skip non-storage mounts
+				if type in MOUNT_IGNORE_FS:
+					continue
+				
+				if any([dir.startswith(d) for d in MOUNT_IGNORE_DIR]):
+					continue
+				
+				if "rw" not in options.split(","):
+					continue
+				
+				# /proc/self/mounts uses a literal \\040 string to escape spaces
+				dir = dir.replace("\\040", " ")
+				
+				basename = os.path.basename(dir)
+				
+				if basename:
+					short_name = " "
+					
+					for c in basename:
+						short_name += c
+						
+						if short_name not in names:
+							break
+				else:
+					short_name = " /"
+				
+				stat = os.statvfs(dir)
+				
+				stor_total = stat.f_blocks * stat.f_bsize
+				stor_free = stat.f_bavail * stat.f_bsize
+				
+				self.fragments.append(short_name)
+				self.fragments.append(style_color(settings.get_space_color(stor_free, stor_total)))
+				stor_free_si = si_number(stor_free)
+				self.fragments.append(space_string(stor_free) %(stor_free_si[0], stor_free_si[1]))
+				self.fragments.append(style_reset())
 
-if __name__ == '__main__':
-	red_thresh = float(sys.argv[1])
-	yellow_thresh = float(sys.argv[2])
-	term_width = int(sys.argv[3])
-	exit_status = int(sys.argv[4])
+class InputPart(Part):
+	"""
+	$
+	or
+	❌42:$
 	
-	lp = LoginPart()
-	pp = PathPart(term_width)
+	$ is either a literal "$" (user) or a "#" (root)
+	"""
+	
+	def __init__(self, settings):
+		Part.__init__(self)
+		self.fragments.append("\n")
+		
+		if settings.return_value:
+			self.fragments.append("%d" %(settings.return_value))
+			self.fragments.append(":")
+		
+		self.fragments.append(style_color(COLOR_PROMPT_ERROR if settings.return_value else COLOR_PROMPT_OK))
+		self.fragments.append("#" if settings.is_root else "$")
+		self.fragments.append(style_reset())
+		self.fragments.append(" ")
+
+if __name__ == "__main__":
+	settings = Settings(sys.argv)
+	sysload = Sysload()
+	
+	lp = LoginPart(settings, sysload)
+	pp = PathPart(settings.term_width)
 	gp = GitPart()
-	pad = Padding(term_width)
-	sp = StatsPart(red_thresh, yellow_thresh)
+	pad = Padding(settings.term_width)
+	sp = StatsPart(settings, sysload)
 	
-	line = [lp, pp, gp, pad, sp]
-	pp.truncate_fit(line)
-	pad.expand_fit(line)
+	top_line = [lp, pp, gp, pad, sp]
+	prompt = InputPart(settings)
 	
-	line_str = ' '.join(str(part) for part in line)
-	sys.stderr.write(line_str)
+	pp.shrink_fit(top_line)
+	pad.expand_fit(top_line)
+	
+	top_line_str = " ".join(str(part) for part in top_line)
+	sys.stderr.write(top_line_str)
+	sys.stderr.write(str(prompt))
 	sys.stderr.flush()
diff --git a/zsh-init b/zsh-init
index c6094fe..4968108 100755
--- a/zsh-init
+++ b/zsh-init
@@ -1,20 +1,11 @@
 #! /bin/zsh
 
-autoload -U colors && colors
-
 OVER_PROMPT_CFG="/etc/over/prompt.cfg"
 
 if [[ -a "$OVER_PROMPT_CFG" ]]; then
 	source "$OVER_PROMPT_CFG"
 fi
 
-function strlen {
-	local PLAIN
-	PLAIN="$(echo $1 | sed -r "s/\x1B\[([0-9]{1,2}(;[0-9]{1,2})?)?[m|K]//g")"
-	
-	echo ${#PLAIN}
-}
-
 function set_title {
 	if [[ ${TERM} == "screen-bce" || ${TERM} == "screen" ]]; then
 		print -Pn "\033k\033${@}\033\134"
@@ -29,55 +20,11 @@ function preexec {
 }
 
 function precmd {
-	local CUT OVER_OPTS RAW_DATA LOGIN_PART STATS_PART DATA TOP_LEFT TOP_RIGHT PADDING PADDING_SIZE GIT_BRANCH COLOR
-	set -A OVER_OPTS ${(s. .)OVER_PROMPT_OPTS}
-	
-	PS1="$(print "%(?.%{\e[1;36m%}.%{\e[1;31m%}%?%{\e[0m%}:%{\e[1;31m%})%(\!.#.$)%{\e[0m%} ")"
-	
-	RAW_DATA="$(/usr/share/over-prompt/data $OVER_OPTS[1] $OVER_OPTS[2] $OVER_OPTS[3])"
-	
-	GIT_BRANCH="$(git rev-parse --abbrev-ref HEAD 2> /dev/null)"
-	
-	if [[ -n "$GIT_BRANCH" ]]; then
-		# rebuild index
-		git update-index -q --ignore-submodules --refresh
-		
-		# so that I can check if there are unstaged changes
-		git diff-files --quiet --ignore-submodules
-		
-		if [[ $? -eq 0 ]]; then
-			COLOR="%{\e[1;32m%}"
-		else
-			COLOR="%{\e[1;31m%}"
-		fi
-		
-		RPS1="$(print "$COLOR$GIT_BRANCH%{\e[0m%}")"
-	else
-		unset RPS1
-	fi
-	
-	if [[ -n "$RAW_DATA" ]]; then
-		set -A DATA ${(s.:::.)RAW_DATA}
-		LOGIN_PART=${DATA[1]}
-		STATS_PART=${DATA[2]}
-		
-		TOP_LEFT=$(print -P "%(\!.\e[1;31m.\e[1;32m)%n\e[0m@$LOGIN_PART")
-		TOP_RIGHT=$(print -P "[ \e[1;36m%T\e[0m | $STATS_PART")
-		PADDING_SIZE=$(($COLUMNS - $(strlen "$TOP_LEFT") - $(strlen "$TOP_RIGHT")))
-		
-# 		if [[ $PADDING_SIZE -lt 0 ]]; then
-# 			CUT=$((0 - $PADDING_SIZE))
-# 			TOP_LEFT="${TOP_LEFT[$CUT,-1]}"
-# 		fi
-		
-		PADDING=$(printf " "%.0s {1..$PADDING_SIZE})
-		
-		print "$TOP_LEFT$PADDING$TOP_RIGHT"
-	else
-		print -P "\e[5;31m!!! unable to run /usr/share/over-prompt/data\e[0m"
-	fi
+	python3 ./render.py $? $COLUMNS 1.25 2.5 0.15 0.05
 	
 	set_title "${PWD}"
 }
 
 unset OVER_PROMPT_CFG
+unset PS1
+unset RPS1