Permalink
Browse files

Host strings now have a 'role' attribute.

This allows to know the 'current' role in a task (with
env.host_string.role), i.e. from which role the current host has been
selected.
  • Loading branch information...
1 parent cabe2b9 commit 47650dc14f406bf56f6210744fd55dca3b7be58f @flupke committed Jan 16, 2013
Showing with 54 additions and 3 deletions.
  1. +11 −3 fabric/task_utils.py
  2. +14 −0 fabric/utils.py
  3. +29 −0 tests/test_main.py
View
@@ -1,4 +1,4 @@
-from fabric.utils import abort, indent
+from fabric.utils import abort, indent, HostString
from fabric import state
@@ -46,11 +46,13 @@ def merge(hosts, roles, exclude, roledefs):
# Look up roles, turn into flat list of hosts
role_hosts = []
+ hosts_roles = [None] * len(list(hosts))
for role in roles:
value = roledefs[role]
# Handle "lazy" roles (callables)
if callable(value):
value = value()
+ hosts_roles += [role] * len(list(value))
role_hosts += value
# Strip whitespace from host strings.
@@ -62,11 +64,17 @@ def merge(hosts, roles, exclude, roledefs):
all_hosts = cleaned_hosts
if state.env.dedupe_hosts:
deduped_hosts = []
- for host in cleaned_hosts:
+ deduped_hosts_roles = []
+ for host, role in zip(cleaned_hosts, hosts_roles):
if host not in deduped_hosts and host not in exclude:
deduped_hosts.append(host)
+ deduped_hosts_roles.append(role)
all_hosts = deduped_hosts
- return all_hosts
+ hosts_roles = deduped_hosts_roles
+
+ # Attach the role attribute to hosts
+ return [HostString(host, role)
+ for host, role in zip(all_hosts, hosts_roles)]
def parse_kwargs(kwargs):
View
@@ -350,3 +350,17 @@ def __setitem__(self, key, value):
raise ValueError("Can't set a slice of a ring buffer!")
else:
return self._super.__setitem__(key, value)
+
+
+class HostString(str):
+ """
+ A subclass of :class:`str` used to store host strings.
+
+ It adds a :attr:`role` attribute, to be able to retrieve which role the
+ host was selected from.
+ """
+
+ def __new__(self, value, role=None):
+ ret = super(HostString, self).__new__(self, value)
+ ret.role = role
+ return ret
View
@@ -312,6 +312,35 @@ def command():
pass
eq_hosts(command, ['a', 'b'], env={'roledefs': lazy_role})
+def test_host_string_role_attr():
+ """
+ Role from which a host was selected can be retrieved with the host
+ string's role attribute.
+ """
+ # Hosts list built from roles
+ fake_env = {'roledefs': {'role1': ['a'], 'role2': ['a']}}
+ hosts = get_hosts(dummy, [], ['role1'], [], fake_env)
+ eq_(hosts, ['a'])
+ eq_(hosts[0].role, 'role1')
+ hosts = get_hosts(dummy, [], ['role2'], [], fake_env)
+ eq_(hosts, ['a'])
+ eq_(hosts[0].role, 'role2')
+ # Hosts list built from hosts
+ hosts = get_hosts(dummy, ['host'], [], [])
+ eq_(hosts, ['host'])
+ eq_(hosts[0].role, None)
+ hosts = get_hosts(dummy, [], [], [], {'hosts': ['host']})
+ eq_(hosts, ['host'])
+ eq_(hosts[0].role, None)
+ # Hosts list built from roles + hosts
+ hosts = get_hosts(dummy, ['b'], ['role1'], [], fake_env)
+ eq_(hosts, ['b', 'a'])
+ eq_(hosts[0].role, None)
+ eq_(hosts[1].role, 'role1')
+ # Hosts list built from roles + hosts (with duplicates)
+ hosts = get_hosts(dummy, ['a'], ['role1'], [], fake_env)
+ eq_(hosts, ['a'])
+ eq_(hosts[0].role, None)
#
# Fabfile loading

0 comments on commit 47650dc

Please sign in to comment.