From eedbc09193336107c441011ad6f4b2a574d5f33f Mon Sep 17 00:00:00 2001
From: Martin Sekera <sekerama@gmail.com>
Date: Sun, 28 Feb 2021 00:32:53 +0100
Subject: [PATCH] fix git command taking too long and crashing

---
 lib/prompt.py | 55 +++++++++++++++++++++++++++++++++------------------
 1 file changed, 36 insertions(+), 19 deletions(-)

diff --git a/lib/prompt.py b/lib/prompt.py
index 6b82abc..bdd01f9 100755
--- a/lib/prompt.py
+++ b/lib/prompt.py
@@ -7,6 +7,7 @@ import pwd
 import socket
 import subprocess
 import sys
+import time
 
 COLOR_LOAD_IDLE = 10
 COLOR_LOAD_OK = 2
@@ -321,14 +322,19 @@ class PathPart(Part):
 
 def command(cmd, timeout=0.5):
 	"""
-	Executes a command, returns stdout, suppresses stderr."
+	Executes a command, returns stdout, suppresses stderr.
+	
+	Kills the process on timeout.
 	"""
 	
 	if type(cmd) == str:
 		cmd = cmd.split()
-	
-	s = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
-	s.wait(timeout)
+	try:
+		s = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+		s.wait(timeout)
+	except subprocess.TimeoutExpired:
+		s.kill()
+		raise
 	
 	return s.stdout.read().decode("utf-8")
 
@@ -410,6 +416,15 @@ class WineprefixPart(Part):
 			else:
 				self.fragments.append(":64")
 
+class TimeBudget:
+	def __init__(self, budget, min=0.1):
+		self.end = time.monotonic() + budget
+		self.min = min
+	
+	@property
+	def trem(self):
+		return max(self.min, self.end - time.monotonic())
+
 class GitPart(Part):
 	"""
 	↘2 ↗4 ⚠M master ?11 ✎6 ✉10
@@ -427,20 +442,21 @@ class GitPart(Part):
 	
 	def __init__(self):
 		Part.__init__(self)
+		tb = TimeBudget(1) # don't spend more than a second here
 		
-		branch_name = command("git name-rev --name-only --no-undefined --always HEAD").strip()
-		
-		if branch_name:
-			try:
+		try:
+			branch_name = command("git name-rev --name-only --no-undefined --always HEAD", tb.trem).strip()
+			
+			if branch_name:
 				self.fragments.append("↱ ")
 				
-				count_to_pull = command("git log --oneline ..@{u}").count("\n")
-				count_to_push = command("git log --oneline @{u}..").count("\n")
-				git_dir = command("git rev-parse --git-dir").strip()
+				count_to_pull = command("git log --oneline ..@{u}", tb.trem).count("\n")
+				count_to_push = command("git log --oneline @{u}..", tb.trem).count("\n")
+				git_dir = command("git rev-parse --git-dir", tb.trem).strip()
 				merging = os.path.exists(os.path.join(git_dir, "MERGE_HEAD"))
-				untracked = command("git ls-files --other --exclude-standard", timeout=0.2).count("\n")
-				modified = command("git diff --name-only").count("\n")
-				staged = command("git diff --name-only --staged").count("\n")
+				untracked = command("git ls-files --other --exclude-standard", tb.trem).count("\n")
+				modified = command("git diff --name-only", tb.trem).count("\n")
+				staged = command("git diff --name-only --staged", tb.trem).count("\n")
 				
 				if count_to_pull:
 					self.fragments.append("↘%d " %(count_to_pull))
@@ -473,11 +489,12 @@ class GitPart(Part):
 					self.fragments.append(" ✉%d" %(staged))
 				
 				self.fragments.append(style_reset())
-			except subprocess.TimeoutExpired:
-				self.fragments = []
-				self.fragments.append(style_color(COLOR_GIT_DIRTY))
-				self.fragments.append("⚠ git timeout ⚠")
-				self.fragments.append(style_reset())
+		
+		except subprocess.TimeoutExpired:
+			self.fragments = []
+			self.fragments.append(style_color(COLOR_GIT_DIRTY))
+			self.fragments.append("⚠ git timeout ⚠")
+			self.fragments.append(style_reset())
 
 class Padding(Part):
 	def __init__(self, term_width):