Refactor cwd segment into a class and add shorten_home argument

powerline/segments/common.py

@@ -82,68 +82,96 @@ 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:
-	:param int dir_shorten_len:
+			return (0, 1, 2)
-		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']())
-	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]"
 		else:
-			raise
+			return super(CwdSegment, self).omitted_args(name, method)
-	home = segment_info['home']
+
-	if home:
+	def get_shortened_path(self, pl, segment_info, shorten_home=True, **kwargs):
-		home = u(home)
+		try:
-		cwd = re.sub('^' + re.escape(home), '~', cwd, 1)
+			path = u(segment_info['getcwd']())
-	cwd_split = cwd.split(os.sep)
+		except OSError as e:
-	cwd_split_len = len(cwd_split)
+			if e.errno == 2:
-	cwd = [i[0:dir_shorten_len] if dir_shorten_len and i else i for i in cwd_split[:-1]] + [cwd_split[-1]]
+				# user most probably deleted the directory
-	if dir_limit_depth and cwd_split_len > dir_limit_depth + 1:
+				# this happens when removing files from Mercurial repos for example
-		del(cwd[0:-dir_limit_depth])
+				pl.warn('Current directory not found')
-		if ellipsis is not None:
+				return "[not found]"
-			cwd.insert(0, ellipsis)
+			else:
-	ret = []
+				raise
-	if not cwd[0]:
+		if shorten_home:
-		cwd[0] = '/'
+			home = segment_info['home']
-	draw_inner_divider = not use_path_separator
+			if home:
-	for part in cwd:
+				home = u(home)
-		if not part:
+				if path.startswith(home):
-			continue
+					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]]
+		if dir_limit_depth and cwd_split_len > dir_limit_depth + 1:
+			del(cwd[0:-dir_limit_depth])
+			if ellipsis is not None:
+				cwd.insert(0, ellipsis)
+		ret = []
+		if not cwd[0]:
+			cwd[0] = '/'
+		draw_inner_divider = not use_path_separator
+		for part in cwd:
+			if not part:
+				continue
+			if use_path_separator:
+				part += os.sep
+			ret.append({
+				'contents': part,
+				'divider_highlight_group': 'cwd:divider',
+				'draw_inner_divider': draw_inner_divider,
+			})
+		ret[-1]['highlight_group'] = ['cwd:current_folder', 'cwd']
 		if use_path_separator:
-			part += os.sep
+			ret[-1]['contents'] = ret[-1]['contents'][:-1]
-		ret.append({
+			if len(ret) > 1 and ret[0]['contents'][0] == os.sep:
-			'contents': part,
+				ret[0]['contents'] = ret[0]['contents'][1:]
-			'divider_highlight_group': 'cwd:divider',
+		return ret
-			'draw_inner_divider': draw_inner_divider,
+
-		})
+
-	ret[-1]['highlight_group'] = ['cwd:current_folder', 'cwd']
+cwd = with_docstring(CwdSegment(),
-	if use_path_separator:
+'''Return the current working directory.
-		ret[-1]['contents'] = ret[-1]['contents'][:-1]
+
-		if len(ret) > 1 and ret[0]['contents'][0] == os.sep:
+Returns a segment list to create a breadcrumb-like effect.
-			ret[0]['contents'] = ret[0]['contents'][1:]
+
-	return ret
+: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):

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']}