Port common.branch segment to KwThreaded*, add common.repository_status

powerline/segments/common.py

@@ -29,25 +29,84 @@ def hostname(only_if_ssh=False):
 	return socket.gethostname()
 
 
-def branch(status_colors=True):
+class RepositorySegment(KwThreadedSegment):
+	def __init__(self):
+		super(RepositorySegment, self).__init__()
+		self.directories = {}
+
+	@staticmethod
+	def key(**kwargs):
+		return os.path.abspath(os.getcwd())
+
+	def update(self):
+		# .compute_state() is running only in this method, and only in one 
+		# thread, thus operations with .directories do not need write locks 
+		# (.render() method is not using .directories). If this is changed 
+		# .directories needs redesigning
+		self.directories.clear()
+		super(RepositorySegment, self).update()
+
+	def compute_state(self, path):
+		repo = guess(path=path)
+		if repo:
+			if repo.directory in self.directories:
+				return self.directories[repo.directory]
+			else:
+				r = self.process_repo(repo)
+				self.directories[repo.directory] = r
+				return r
+
+
+class RepositoryStatusSegment(RepositorySegment):
+	interval = 2
+
+	@staticmethod
+	def process_repo(repo):
+		return repo.status()
+
+
+repository_status = with_docstring(RepositoryStatusSegment(),
+'''Return the status for the current VCS repository.''')
+
+
+class BranchSegment(RepositorySegment):
+	interval = 0.2
+	started_repository_status = False
+
+	@staticmethod
+	def process_repo(repo):
+		return repo.branch()
+
+	@staticmethod
+	def render_one(branch, status_colors=False, **kwargs):
+		if branch and status_colors:
+			return [{
+				'contents': branch,
+				'highlight_group': ['branch_dirty' if repository_status() else 'branch_clean', 'branch'],
+				}]
+		else:
+			return branch
+
+	def startup(self, status_colors=False, **kwargs):
+		super(BranchSegment, self).startup()
+		if status_colors:
+			self.started_repository_status = True
+			repository_status.startup()
+
+	def shutdown(self):
+		if self.started_repository_status:
+			repository_status.shutdown()
+		super(BranchSegment, self).shutdown()
+
+
+branch = with_docstring(BranchSegment(),
 '''Return the current VCS branch.
 
 :param bool status_colors:
 	determines whether repository status will be used to determine highlighting. Default: True.
 
 Highlight groups used: ``branch_clean``, ``branch_dirty``, ``branch``.
-	'''
+''')
-	repo = guess(path=os.path.abspath(os.getcwd()))
-	if repo:
-		branch = repo.branch()
-		if status_colors:
-			return [{
-				'contents': branch,
-				'highlight_group': ['branch_dirty' if repo.status() else 'branch_clean', 'branch'],
-				}]
-		else:
-			return branch
-	return None
 
 
 def cwd(dir_shorten_len=None, dir_limit_depth=None):

tests/test_segments.py

@@ -48,14 +48,15 @@ class TestCommon(TestCase):
 				self.assertEqual(common.user(), [{'contents': 'def', 'highlight_group': ['superuser', 'user']}])
 
 	def test_branch(self):
-		with replace_module_attr(common, 'guess', lambda path: Args(branch=lambda: os.path.basename(path), status=lambda: None)):
+		with replace_module_attr(common, 'guess', lambda path: Args(branch=lambda: os.path.basename(path), status=lambda: None, directory='/tmp/tests')):
 			self.assertEqual(common.branch(status_colors=False), 'tests')
 			self.assertEqual(common.branch(status_colors=True),
 					[{'contents': 'tests', 'highlight_group': ['branch_clean', 'branch']}])
-		with replace_module_attr(common, 'guess', lambda path: Args(branch=lambda: os.path.basename(path), status=lambda: 'D  ')):
+		with replace_module_attr(common, 'guess', lambda path: Args(branch=lambda: os.path.basename(path), status=lambda: 'D  ', directory='/tmp/tests')):
 			self.assertEqual(common.branch(status_colors=False), 'tests')
-			self.assertEqual(common.branch(),
+			self.assertEqual(common.branch(status_colors=True),
 					[{'contents': 'tests', 'highlight_group': ['branch_dirty', 'branch']}])
+			self.assertEqual(common.branch(), 'tests')
 		with replace_module_attr(common, 'guess', lambda path: None):
 			self.assertEqual(common.branch(), None)