Refactor cwd segment into a class and add shorten_home argument

powerline/segments/common.py

@@ -82,41 +82,44 @@ def branch(pl, segment_info, create_watcher, status_colors=False):
 
 
 @requires_segment_info
-def cwd(pl, segment_info, dir_shorten_len=None, dir_limit_depth=None, use_path_separator=False, ellipsis='⋯'):
+class CwdSegment(Segment):
-	'''Return the current working directory.
+	def argspecobjs(self):
+		for obj in super(CwdSegment, self).argspecobjs():
+			yield obj
+		yield 'get_shortened_path', self.get_shortened_path
 
-	Returns a segment list to create a breadcrumb-like effect.
+	def omitted_args(self, name, method):
+		if method is self.get_shortened_path:
+			return (0, 1, 2)
+		else:
+			return super(CwdSegment, self).omitted_args(name, method)
 
-	:param int dir_shorten_len:
+	def get_shortened_path(self, pl, segment_info, shorten_home=True, **kwargs):
-		shorten parent directory names to this length (e.g. 
-		:file:`/long/path/to/powerline` → :file:`/l/p/t/powerline`)
-	:param int dir_limit_depth:
-		limit directory depth to this number (e.g. 
-		:file:`/long/path/to/powerline` → :file:`⋯/to/powerline`)
-	:param bool use_path_separator:
-		Use path separator in place of soft divider.
-	:param str ellipsis:
-		Specifies what to use in place of omitted directories. Use None to not 
-		show this subsegment at all.
-
-	Divider highlight group used: ``cwd:divider``.
-
-	Highlight groups used: ``cwd:current_folder`` or ``cwd``. It is recommended to define all highlight groups.
-	'''
 		try:
-		cwd = u(segment_info['getcwd']())
+			path = u(segment_info['getcwd']())
 		except OSError as e:
 			if e.errno == 2:
 				# user most probably deleted the directory
 				# this happens when removing files from Mercurial repos for example
 				pl.warn('Current directory not found')
-			cwd = "[not found]"
+				return "[not found]"
 			else:
 				raise
+		if shorten_home:
 			home = segment_info['home']
 			if home:
 				home = u(home)
-		cwd = re.sub('^' + re.escape(home), '~', cwd, 1)
+				if path.startswith(home):
+					path = '~' + path[len(home):]
+		return path
+
+	def __call__(self, pl, segment_info,
+	             dir_shorten_len=None,
+	             dir_limit_depth=None,
+	             use_path_separator=False,
+	             ellipsis='⋯',
+	             **kwargs):
+		cwd = self.get_shortened_path(pl, segment_info, **kwargs)
 		cwd_split = cwd.split(os.sep)
 		cwd_split_len = len(cwd_split)
 		cwd = [i[0:dir_shorten_len] if dir_shorten_len and i else i for i in cwd_split[:-1]] + [cwd_split[-1]]
@@ -146,6 +149,31 @@ def cwd(pl, segment_info, dir_shorten_len=None, dir_limit_depth=None, use_path_s
 		return ret
 
 
+cwd = with_docstring(CwdSegment(),
+'''Return the current working directory.
+
+Returns a segment list to create a breadcrumb-like effect.
+
+:param int dir_shorten_len:
+	shorten parent directory names to this length (e.g. 
+	:file:`/long/path/to/powerline` → :file:`/l/p/t/powerline`)
+:param int dir_limit_depth:
+	limit directory depth to this number (e.g. 
+	:file:`/long/path/to/powerline` → :file:`⋯/to/powerline`)
+:param bool use_path_separator:
+	Use path separator in place of soft divider.
+:param bool shorten_home:
+	Shorten home directory to ``~``.
+:param str ellipsis:
+	Specifies what to use in place of omitted directories. Use None to not 
+	show this subsegment at all.
+
+Divider highlight group used: ``cwd:divider``.
+
+Highlight groups used: ``cwd:current_folder`` or ``cwd``. It is recommended to define all highlight groups.
+''')
+
+
 def date(pl, format='%Y-%m-%d', istime=False):
 	'''Return the current date.
 

tests/test_segments.py

@@ -277,6 +277,12 @@ class TestCommon(TestCase):
 				{'contents': 'foo', 'divider_highlight_group': 'cwd:divider', 'draw_inner_divider': True},
 				{'contents': 'bar', 'divider_highlight_group': 'cwd:divider', 'draw_inner_divider': True, 'highlight_group': ['cwd:current_folder', 'cwd']}
 			])
+			self.assertEqual(common.cwd(pl=pl, segment_info=segment_info, dir_limit_depth=3, shorten_home=False), [
+				{'contents': '⋯', 'divider_highlight_group': 'cwd:divider', 'draw_inner_divider': True},
+				{'contents': 'ghi', 'divider_highlight_group': 'cwd:divider', 'draw_inner_divider': True},
+				{'contents': 'foo', 'divider_highlight_group': 'cwd:divider', 'draw_inner_divider': True},
+				{'contents': 'bar', 'divider_highlight_group': 'cwd:divider', 'draw_inner_divider': True, 'highlight_group': ['cwd:current_folder', 'cwd']}
+			])
 			self.assertEqual(common.cwd(pl=pl, segment_info=segment_info, dir_limit_depth=1), [
 				{'contents': '⋯', 'divider_highlight_group': 'cwd:divider', 'draw_inner_divider': True},
 				{'contents': 'bar', 'divider_highlight_group': 'cwd:divider', 'draw_inner_divider': True, 'highlight_group': ['cwd:current_folder', 'cwd']}