From 918dd957cf8efc66a849fa0cf29e46992bde2d9b Mon Sep 17 00:00:00 2001
From: S0lll0s <s0lll0s@blinkenshell.org>
Date: Fri, 21 Aug 2015 17:50:27 +0200
Subject: [PATCH] Add 'output' option to segments.i3wm.workspaces

---
 powerline/segments/i3wm.py | 10 ++++++++--
 tests/test_segments.py     | 17 +++++++++++++----
 2 files changed, 21 insertions(+), 6 deletions(-)

diff --git a/powerline/segments/i3wm.py b/powerline/segments/i3wm.py
index f0c16427..a5a82774 100644
--- a/powerline/segments/i3wm.py
+++ b/powerline/segments/i3wm.py
@@ -19,7 +19,7 @@ def calcgrp(w):
 	return group
 
 
-def workspaces(pl, only_show=None, strip=0):
+def workspaces(pl, only_show=None, output=None, strip=0):
 	'''Return list of used workspaces
 
 	:param list only_show:
@@ -27,6 +27,9 @@ def workspaces(pl, only_show=None, strip=0):
 		``"urgent"`` and ``"focused"``. If omitted or ``null`` all workspaces 
 		are shown.
 
+	:param str output:
+		If specified, only workspaces on this output are shown.
+
 	:param int strip:
 		Specifies how many characters from the front of each workspace name 
 		should be stripped (e.g. to remove workspace numbers). Defaults to zero.
@@ -45,7 +48,10 @@ def workspaces(pl, only_show=None, strip=0):
 	return [{
 		'contents': w['name'][min(len(w['name']), strip):],
 		'highlight_groups': calcgrp(w)
-	} for w in conn.get_workspaces() if not only_show or any((w[typ] for typ in only_show))]
+	} for w in conn.get_workspaces()
+		if (not only_show or any(w[typ] for typ in only_show))
+		and (not output or w['output'] == output)
+	]
 
 
 @requires_segment_info
diff --git a/tests/test_segments.py b/tests/test_segments.py
index c01dd6d5..b3c77029 100644
--- a/tests/test_segments.py
+++ b/tests/test_segments.py
@@ -819,10 +819,10 @@ class TestI3WM(TestCase):
 	def test_workspaces(self):
 		pl = Pl()
 		with replace_attr(i3wm, 'conn', Args(get_workspaces=lambda: iter([
-			{'name': '1: w1', 'focused': False, 'urgent': False, 'visible': False},
-			{'name': '2: w2', 'focused': False, 'urgent': False, 'visible': True},
-			{'name': '3: w3', 'focused': False, 'urgent': True, 'visible': True},
-			{'name': '4: w4', 'focused': True, 'urgent': True, 'visible': True},
+			{'name': '1: w1', 'output': 'LVDS1', 'focused': False, 'urgent': False, 'visible': False},
+			{'name': '2: w2', 'output': 'LVDS1', 'focused': False, 'urgent': False, 'visible': True},
+			{'name': '3: w3', 'output': 'HDMI1', 'focused': False, 'urgent': True, 'visible': True},
+			{'name': '4: w4', 'output': 'DVI01', 'focused': True, 'urgent': True, 'visible': True},
 		]))):
 			self.assertEqual(i3wm.workspaces(pl=pl), [
 				{'contents': '1: w1', 'highlight_groups': ['workspace']},
@@ -850,6 +850,15 @@ class TestI3WM(TestCase):
 				{'contents': 'w3', 'highlight_groups': ['w_urgent', 'w_visible', 'workspace']},
 				{'contents': 'w4', 'highlight_groups': ['w_focused', 'w_urgent', 'w_visible', 'workspace']},
 			])
+			self.assertEqual(i3wm.workspaces(pl=pl, only_show=['focused', 'urgent'], output='DVI01'), [
+				{'contents': '4: w4', 'highlight_groups': ['w_focused', 'w_urgent', 'w_visible', 'workspace']},
+			])
+			self.assertEqual(i3wm.workspaces(pl=pl, only_show=['visible'], output='HDMI1'), [
+				{'contents': '3: w3', 'highlight_groups': ['w_urgent', 'w_visible', 'workspace']},
+			])
+			self.assertEqual(i3wm.workspaces(pl=pl, only_show=['visible'], strip=3, output='LVDS1'), [
+				{'contents': 'w2', 'highlight_groups': ['w_visible', 'workspace']},
+			])
 
 	def test_mode(self):
 		pl = Pl()