Refactor Vim and common branch segments to share code

powerline/segments/common/vcs.py

@@ -2,32 +2,45 @@
 from __future__ import (unicode_literals, division, absolute_import, print_function)
 
 from powerline.lib.vcs import guess, tree_status
+from powerline.segments import Segment, with_docstring
 from powerline.theme import requires_segment_info, requires_filesystem_watcher
 
 
 @requires_filesystem_watcher
 @requires_segment_info
-def branch(pl, segment_info, create_watcher, status_colors=False):
-	'''Return the current VCS branch.
+class BranchSegment(Segment):
+	divider_highlight_group = None
 
-	:param bool status_colors:
-		determines whether repository status will be used to determine highlighting. Default: False.
+	@staticmethod
+	def get_directory(segment_info):
+		return segment_info['getcwd']()
 
-	Highlight groups used: ``branch_clean``, ``branch_dirty``, ``branch``.
-	'''
-	name = segment_info['getcwd']()
-	repo = guess(path=name, create_watcher=create_watcher)
-	if repo is not None:
-		branch = repo.branch()
-		scol = ['branch']
-		if status_colors:
-			try:
-				status = tree_status(repo, pl)
-			except Exception as e:
-				pl.exception('Failed to compute tree status: {0}', str(e))
-				status = '?'
-			scol.insert(0, 'branch_dirty' if status and status.strip() else 'branch_clean')
-		return [{
-			'contents': branch,
-			'highlight_group': scol,
-		}]
+	def __call__(self, pl, segment_info, create_watcher, status_colors=False):
+		name = self.get_directory(segment_info)
+		if name:
+			repo = guess(path=name, create_watcher=create_watcher)
+			if repo is not None:
+				branch = repo.branch()
+				scol = ['branch']
+				if status_colors:
+					try:
+						status = tree_status(repo, pl)
+					except Exception as e:
+						pl.exception('Failed to compute tree status: {0}', str(e))
+						status = '?'
+					scol.insert(0, 'branch_dirty' if status and status.strip() else 'branch_clean')
+				return [{
+					'contents': branch,
+					'highlight_group': scol,
+					'divider_highlight_group': self.divider_highlight_group,
+				}]
+
+
+branch = with_docstring(BranchSegment(),
+'''Return the current VCS branch.
+
+:param bool status_colors:
+	determines whether repository status will be used to determine highlighting. Default: False.
+
+Highlight groups used: ``branch_clean``, ``branch_dirty``, ``branch``.
+''')

powerline/segments/vim/__init__.py

@@ -17,9 +17,11 @@ from powerline.bindings.vim import (vim_get_func, getbufvar, vim_getbufoption,
                                     list_tabpage_buffers_segment_info)
 from powerline.theme import requires_segment_info, requires_filesystem_watcher
 from powerline.lib import add_divider_highlight_group
-from powerline.lib.vcs import guess, tree_status
+from powerline.lib.vcs import guess
 from powerline.lib.humanize_bytes import humanize_bytes
 from powerline.lib import wraps_saveargs as wraps
+from powerline.segments.common.vcs import BranchSegment
+from powerline.segments import with_docstring
 
 try:
 	from __builtin__ import xrange as range
@@ -480,31 +482,26 @@ def modified_buffers(pl, text='+ ', join_str=','):
 
 @requires_filesystem_watcher
 @requires_segment_info
-def branch(pl, segment_info, create_watcher, status_colors=False):
-	'''Return the current working branch.
+class VimBranchSegment(BranchSegment):
+	divider_highlight_group = 'branch:divider'
 
-	:param bool status_colors:
-		determines whether repository status will be used to determine highlighting. Default: False.
+	@staticmethod
+	def get_directory(segment_info):
+		if vim_getbufoption(segment_info, 'buftype'):
+			return None
+		return buffer_name(segment_info)
 
-	Highlight groups used: ``branch_clean``, ``branch_dirty``, ``branch``.
 
-	Divider highlight group used: ``branch:divider``.
-	'''
-	name = buffer_name(segment_info)
-	skip = not (name and (not vim_getbufoption(segment_info, 'buftype')))
-	if not skip:
-		repo = guess(path=name, create_watcher=create_watcher)
-		if repo is not None:
-			branch = repo.branch()
-			scol = ['branch']
-			if status_colors:
-				status = tree_status(repo, pl)
-				scol.insert(0, 'branch_dirty' if status and status.strip() else 'branch_clean')
-			return [{
-				'contents': branch,
-				'highlight_group': scol,
-				'divider_highlight_group': 'branch:divider',
-			}]
+branch = with_docstring(VimBranchSegment(),
+'''Return the current working branch.
+
+:param bool status_colors:
+	determines whether repository status will be used to determine highlighting. Default: False.
+
+Highlight groups used: ``branch_clean``, ``branch_dirty``, ``branch``.
+
+Divider highlight group used: ``branch:divider``.
+''')
 
 
 @requires_filesystem_watcher

tests/test_segments.py

@@ -600,23 +600,33 @@ class TestVcs(TestCommon):
 		branch = partial(common.branch, pl=pl, create_watcher=create_watcher)
 		with replace_attr(self.module, 'guess', get_dummy_guess(status=lambda: None, directory='/tmp/tests')):
 			with replace_attr(self.module, 'tree_status', lambda repo, pl: None):
-				self.assertEqual(branch(segment_info=segment_info, status_colors=False), [
-					{'highlight_group': ['branch'], 'contents': 'tests'}
-				])
-				self.assertEqual(branch(segment_info=segment_info, status_colors=True), [
-					{'contents': 'tests', 'highlight_group': ['branch_clean', 'branch']}
-				])
+				self.assertEqual(branch(segment_info=segment_info, status_colors=False), [{
+						'highlight_group': ['branch'],
+						'contents': 'tests',
+						'divider_highlight_group': None
+				}])
+				self.assertEqual(branch(segment_info=segment_info, status_colors=True), [{
+					'contents': 'tests',
+					'highlight_group': ['branch_clean', 'branch'],
+					'divider_highlight_group': None
+				}])
 		with replace_attr(self.module, 'guess', get_dummy_guess(status=lambda: 'D  ', directory='/tmp/tests')):
 			with replace_attr(self.module, 'tree_status', lambda repo, pl: 'D '):
-				self.assertEqual(branch(segment_info=segment_info, status_colors=False), [
-					{'highlight_group': ['branch'], 'contents': 'tests'}
-				])
-				self.assertEqual(branch(segment_info=segment_info, status_colors=True), [
-					{'contents': 'tests', 'highlight_group': ['branch_dirty', 'branch']}
-				])
-				self.assertEqual(branch(segment_info=segment_info, status_colors=False), [
-					{'highlight_group': ['branch'], 'contents': 'tests'}
-				])
+				self.assertEqual(branch(segment_info=segment_info, status_colors=False), [{
+					'highlight_group': ['branch'],
+					'contents': 'tests',
+					'divider_highlight_group': None
+				}])
+				self.assertEqual(branch(segment_info=segment_info, status_colors=True), [{
+					'contents': 'tests',
+					'highlight_group': ['branch_dirty', 'branch'],
+					'divider_highlight_group': None
+				}])
+				self.assertEqual(branch(segment_info=segment_info, status_colors=False), [{
+					'highlight_group': ['branch'],
+					'contents': 'tests',
+					'divider_highlight_group': None
+				}])
 		with replace_attr(self.module, 'guess', lambda path, create_watcher: None):
 			self.assertEqual(branch(segment_info=segment_info, status_colors=False), None)
 
@@ -1057,16 +1067,16 @@ class TestVim(TestCase):
 		create_watcher = get_fallback_create_watcher()
 		branch = partial(self.vim.branch, pl=pl, create_watcher=create_watcher)
 		with vim_module._with('buffer', '/foo') as segment_info:
-			with replace_attr(self.vim, 'guess', get_dummy_guess(status=lambda: None)):
-				with replace_attr(self.vim, 'tree_status', lambda repo, pl: None):
+			with replace_attr(self.vcs, 'guess', get_dummy_guess(status=lambda: None)):
+				with replace_attr(self.vcs, 'tree_status', lambda repo, pl: None):
 					self.assertEqual(branch(segment_info=segment_info, status_colors=False), [
 						{'divider_highlight_group': 'branch:divider', 'highlight_group': ['branch'], 'contents': 'foo'}
 					])
 					self.assertEqual(branch(segment_info=segment_info, status_colors=True), [
 						{'divider_highlight_group': 'branch:divider', 'highlight_group': ['branch_clean', 'branch'], 'contents': 'foo'}
 					])
-			with replace_attr(self.vim, 'guess', get_dummy_guess(status=lambda: 'DU')):
-				with replace_attr(self.vim, 'tree_status', lambda repo, pl: 'DU'):
+			with replace_attr(self.vcs, 'guess', get_dummy_guess(status=lambda: 'DU')):
+				with replace_attr(self.vcs, 'tree_status', lambda repo, pl: 'DU'):
 					self.assertEqual(branch(segment_info=segment_info, status_colors=False), [
 						{'divider_highlight_group': 'branch:divider', 'highlight_group': ['branch'], 'contents': 'foo'}
 					])
@@ -1155,6 +1165,8 @@ class TestVim(TestCase):
 		sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), 'path')))
 		from powerline.segments import vim
 		cls.vim = vim
+		from powerline.segments.common import vcs
+		cls.vcs = vcs
 
 	@classmethod
 	def tearDownClass(cls):