From 79e2f2aee62634ff5495da97ce0033baa3f3c7eb Mon Sep 17 00:00:00 2001
From: ZyX <kp-pav@ya.ru>
Date: Thu, 13 Dec 2012 22:13:03 +0400
Subject: [PATCH] Add a way to import extension segments from more convenient
 locations

I.e. not powerline.ext.{ext}.segments.{smth}, but just {smth}.

Refs #3.
---
 powerline/segments.py | 58 +++++++++++++++++++++++++++++++++++++++++++
 powerline/theme.py    | 42 +++----------------------------
 2 files changed, 62 insertions(+), 38 deletions(-)
 create mode 100644 powerline/segments.py

diff --git a/powerline/segments.py b/powerline/segments.py
new file mode 100644
index 00000000..665dcdf7
--- /dev/null
+++ b/powerline/segments.py
@@ -0,0 +1,58 @@
+# -*- coding: utf-8 -*-
+
+from importlib import import_module
+import sys
+import os
+
+
+class Segments(object):
+	def __init__(self, ext, common_config, colorscheme):
+		self.ext = ext
+		self.path = [os.path.expanduser(path) for path in common_config.get('paths', [])]
+		self.colorscheme = colorscheme
+
+	def get_function(self, segment):
+		oldpath = sys.path
+		sys.path = self.path + sys.path
+		segment_module = str(segment.get('module', 'powerline.ext.{0}.segments.core'.format(self.ext)))
+
+		try:
+			return None, getattr(import_module(segment_module), segment['name']), segment_module
+		finally:
+			sys.path = oldpath
+
+	@staticmethod
+	def get_string(segment):
+		return segment.get('contents'), None, None
+
+	@staticmethod
+	def get_filler(segment):
+		return None, None, None
+
+	def get(self, segment, side):
+		segment_type = segment.get('type', 'function')
+
+		try:
+			contents, contents_func, key = getattr(self, 'get_{0}'.format(segment_type))(segment)
+		except AttributeError:
+			raise TypeError('Unknown segment type: {0}'.format(segment_type))
+
+		highlighting_group = segment.get('highlight', segment.get('name'))
+
+		return {
+			'key': key,
+			'type': segment_type,
+			'highlight': self.colorscheme.get_group_highlighting(highlighting_group),
+			'before': segment.get('before', ''),
+			'after': segment.get('after', ''),
+			'contents_func': contents_func,
+			'contents': contents,
+			'args': segment.get('args', {}),
+			'ljust': segment.get('ljust', False),
+			'rjust': segment.get('rjust', False),
+			'priority': segment.get('priority', -1),
+			'draw_divider': segment.get('draw_divider', True),
+			'side': side,
+			'exclude_modes': segment.get('exclude_modes', []),
+			'include_modes': segment.get('include_modes', []),
+		}
diff --git a/powerline/theme.py b/powerline/theme.py
index 9ae869d9..e52b7260 100644
--- a/powerline/theme.py
+++ b/powerline/theme.py
@@ -1,6 +1,6 @@
 # -*- coding: utf-8 -*-
 
-import importlib
+from powerline.segments import Segments
 
 
 class Theme(object):
@@ -14,44 +14,10 @@ class Theme(object):
 			'highlight': {self.colorscheme.DEFAULT_MODE_KEY: {'fg': (False, False), 'bg': (False, False), 'attr': 0}}
 		}
 
+		get_segment = Segments(ext, common_config, colorscheme).get
+
 		for side in ['left', 'right']:
-			for segment in theme_config['segments'].get(side, []):
-				contents = None
-				contents_func = None
-				segment_type = segment.get('type', 'function')
-				segment_module = segment.get('module', 'core')
-
-				if segment_type == 'function':
-					# Import segment function and assign it to the contents
-					function_module = 'powerline.ext.{0}.segments.{1}'.format(ext, segment_module)
-					function_name = segment['name']
-					contents_func = getattr(importlib.import_module(function_module), function_name)
-				elif segment_type == 'string':
-					contents = segment.get('contents')
-				elif segment_type == 'filler':
-					pass
-				else:
-					raise TypeError('Unknown segment type: {0}'.format(segment_type))
-
-				highlighting_group = segment.get('highlight', segment.get('name'))
-
-				self.segments.append({
-					'key': None if segment_type != 'function' else '{0}.{1}'.format(segment_module, function_name),
-					'type': segment_type,
-					'highlight': self.colorscheme.get_group_highlighting(highlighting_group),
-					'before': segment.get('before', ''),
-					'after': segment.get('after', ''),
-					'contents_func': contents_func,
-					'contents': contents,
-					'args': segment.get('args', {}),
-					'ljust': segment.get('ljust', False),
-					'rjust': segment.get('rjust', False),
-					'priority': segment.get('priority', -1),
-					'draw_divider': segment.get('draw_divider', True),
-					'side': side,
-					'exclude_modes': segment.get('exclude_modes', []),
-					'include_modes': segment.get('include_modes', []),
-				})
+			self.segments.extend((get_segment(segment, side) for segment in theme_config['segments'].get(side, [])))
 
 	def get_divider(self, side='left', type='soft'):
 		'''Return segment divider.