Refactor watch context manager

8 files changed, 89 insertions, 51 deletions

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 0f7eb7f..bf5f728 100644
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -4,6 +4,9 @@ Changelog
 Version 0.5 (unreleased)
 ------------------------
 
+* The ``watch`` context manager now allows you to either provide
+  a callback or do an explicit check afterwards (**warning**: this change
+  is not backwards compatible, please update your fabfiles)
 * Add support for some network-related operations:
     * get the IPV4 address assigned to an interface
     * get the list of name server IP addresses
diff --git a/fabtools/files.py b/fabtools/files.py
index 7463b73..9177ad6 100644
--- a/fabtools/files.py
+++ b/fabtools/files.py
@@ -4,7 +4,6 @@ Fabric tools for managing files and directories
 from __future__ import with_statement
 
 import os.path
-from contextlib import contextmanager
 
 from fabric.api import *
 from fabric.contrib.files import upload_template as _upload_template
@@ -67,31 +66,61 @@ def md5sum(filename, use_sudo=False):
     return res.split()[0]
 
 
-@contextmanager
-def watch(_filenames, _use_sudo=False, _callable=None, *args, **kwargs):
+class watch(object):
     """
-    Trigger the callable if any files changed at the end of the context.
-    Underscores are used to avoid conflicts with the args/kwargs
-    of the callable.
-    The filenames can be a string (short for [filename1]) instead of a list.
+    Context manager to watch for changes to the contents of some files
 
-    Typical usage:
+    The 'filenames' argument can be either a string (single filename)
+    or a list (multiple filenames).
+
+    You can read the 'changed' attribute at the end of the block to
+    check if the contents of any of the watched files has changed.
+
+    You can also provide a callback that will be called at the end of
+    the block if the contents of any of the watched files has changed.
+
+    Example using an explicit check::
+
+        from fabric.contrib.files import comment, uncomment
+        from fabtools.files import watch
+        from fabtools.services import restart
+        with watch('/etc/daemon.conf') as config:
+            uncomment('/etc/daemon.conf', 'someoption')
+            comment('/etc/daemon.conf', 'otheroption')
+        if config.changed:
+            restart('daemon')
+
+    Example using a callback::
+
+        from functools import partial
+        from fabric.contrib.files import comment, uncomment
         from fabtools.files import watch
         from fabtools.services import restart
-        from fabtools.contrib.files import uncomment
-        with watch("/etc/daemon.conf", True, restart, "daemon"):
-            uncomment("/etc/daemon.conf", "someoption")
-            comment("/etc/daemon.conf", "otheroption")
-        # The daemon will be restarted only if its config file was changed.
+        with watch('/etc/daemon.conf', callback=partial(restart, 'daemon')):
+            uncomment('/etc/daemon.conf', 'someoption')
+            comment('/etc/daemon.conf', 'otheroption')
     """
-    filenames = [_filenames] if isinstance(_filenames, basestring) \
-                else _filenames
-    old_md5 = dict()
-    with settings(hide('warnings')):
-        for filename in filenames:
-            old_md5[filename] = md5sum(filename, _use_sudo)
-    yield
-    for filename in filenames:
-        if md5sum(filename, _use_sudo) != old_md5[filename]:
-            _callable(*args, **kwargs)
-            return
+
+    def __init__(self, filenames, callback=None, use_sudo=False):
+        if isinstance(filenames, basestring):
+            self.filenames = [filenames]
+        else:
+            self.filenames = filenames
+        self.callback = callback
+        self.use_sudo = use_sudo
+        self.digest = dict()
+        self.changed = False
+
+    def __enter__(self):
+        with settings(hide('warnings')):
+            for filename in self.filenames:
+                self.digest[filename] = md5sum(filename, self.use_sudo)
+        return self
+
+    def __exit__(self, type, value, tb):
+        for filename in self.filenames:
+            if md5sum(filename, self.use_sudo) != self.digest[filename]:
+                self.changed = True
+                break
+        if self.changed and self.callback:
+            self.callback()
diff --git a/fabtools/require/deb.py b/fabtools/require/deb.py
index a26e15f..933aaf4 100644
--- a/fabtools/require/deb.py
+++ b/fabtools/require/deb.py
@@ -17,11 +17,11 @@ def source(name, uri, distribution, *components):
     path = '/etc/apt/sources.list.d/%(name)s.list' % locals()
     components = ' '.join(components)
     source_line = 'deb %(uri)s %(distribution)s %(components)s\n' % locals()
-    def on_update():
+    with watch(path) as config:
+        fabtools.require.file(path=path, contents=source_line, use_sudo=True)
+    if config.changed:
         puts('Added APT repository: %s' % source_line)
         update_index()
-    with watch(path, _callable=on_update):
-        fabtools.require.file(path=path, contents=source_line, use_sudo=True)
 
 
 def ppa(name):
diff --git a/fabtools/require/redis.py b/fabtools/require/redis.py
index 3238231..cc9041c 100644
--- a/fabtools/require/redis.py
+++ b/fabtools/require/redis.py
@@ -91,12 +91,7 @@ def instance(name, version=VERSION, **kwargs):
     config_filename = '/etc/redis/%(name)s.conf' % locals()
 
     # Upload config file
-    context = dict(need_restart=False)
-
-    def on_change():
-        context['need_restart'] = True
-
-    with watch(config_filename, True, on_change):
+    with watch(config_filename, use_sudo=True) as config:
         require.file(config_filename, contents='\n'.join(lines),
             use_sudo=True, owner='redis')
 
@@ -106,5 +101,7 @@ def instance(name, version=VERSION, **kwargs):
         user='redis',
         directory='/var/run/redis',
         command="%(redis_server)s %(config_filename)s" % locals())
-    if context['need_restart']:
+
+    # Restart if needed
+    if config.changed:
         fabtools.supervisor.restart_process(process_name)
diff --git a/fabtools/require/shorewall.py b/fabtools/require/shorewall.py
index 9e19e58..51830d8 100644
--- a/fabtools/require/shorewall.py
+++ b/fabtools/require/shorewall.py
@@ -249,12 +249,7 @@ def firewall(zones=None, interfaces=None, policy=None, rules=None,
     """
     package('shorewall')
 
-    def on_change():
-        puts("Shorewall configuration changed")
-        if is_started():
-            restart('shorewall')
-
-    with watch(CONFIG_FILES, False, on_change):
+    with watch(CONFIG_FILES) as config:
         _zone_config(zones)
         _interfaces_config(interfaces)
         _policy_config(policy)
@@ -262,6 +257,11 @@ def firewall(zones=None, interfaces=None, policy=None, rules=None,
         _routestopped_config(routestopped)
         _masq_config(masq)
 
+    if config.changed:
+        puts("Shorewall configuration changed")
+        if is_started():
+            restart('shorewall')
+
     with settings(hide('running')):
         sed('/etc/default/shorewall', 'startup=0', 'startup=1', use_sudo=True)
 
diff --git a/fabtools/require/supervisor.py b/fabtools/require/supervisor.py
index 28955e5..07b6553 100644
--- a/fabtools/require/supervisor.py
+++ b/fabtools/require/supervisor.py
@@ -29,8 +29,7 @@ def process(name, **kwargs):
 
     # Upload config file
     filename = '/etc/supervisor/conf.d/%(name)s.conf' % locals()
-
-    with watch(filename, True, update_config):
+    with watch(filename, callback=update_config, use_sudo=True):
         require.file(filename, contents='\n'.join(lines), use_sudo=True)
 
     # Start the process if needed
diff --git a/fabtools/require/system.py b/fabtools/require/system.py
index 3731251..d798fbd 100644
--- a/fabtools/require/system.py
+++ b/fabtools/require/system.py
@@ -24,12 +24,12 @@ def sysctl(key, value, persist=True):
     if persist:
         from fabtools import require
         filename = '/etc/sysctl.d/60-%s.conf' % key
-        def on_change():
-            sudo('service procps start')
-        with watch(filename, True, on_change):
+        with watch(filename, use_sudo=True) as config:
             require.file(filename,
                 contents='%(key)s = %(value)s\n' % locals(),
                 use_sudo=True)
+        if config.changed:
+            sudo('service procps start')
 
 
 def hostname(name):
@@ -47,11 +47,8 @@ def locales(names):
 
     config_file = '/var/lib/locales/supported.d/local'
 
-    def regenerate():
-        sudo('dpkg-reconfigure locales')
-
     # Regenerate locales if config file changes
-    with watch(config_file, True, regenerate):
+    with watch(config_file, use_sudo=True) as config:
 
         # Add valid locale names to the config file
         supported = dict(supported_locales())
@@ -63,6 +60,9 @@ def locales(names):
             else:
                 warn('Unsupported locale name "%s"' % name)
 
+    if config.changed:
+        sudo('dpkg-reconfigure locales')
+
 
 def locale(name):
     """
diff --git a/fabtools/tests/fabfiles/files.py b/fabtools/tests/fabfiles/files.py
index 1832792..53fe462 100644
--- a/fabtools/tests/fabfiles/files.py
+++ b/fabtools/tests/fabfiles/files.py
@@ -2,6 +2,7 @@ from __future__ import with_statement
 
 import os
 from tempfile import mkstemp
+from functools import partial
 
 from fabric.api import *
 from fabtools import require
@@ -40,12 +41,21 @@ def files():
         assert fabtools.files.is_file('baz')
         assert run('cat baz') == baz_contents, run('cat baz')
 
+        # Ensure that changes to watched file are detected
+        require.file('watched', contents='aaa')
+        with fabtools.files.watch('watched') as f:
+            require.file('watched', contents='bbb')
+        assert f.changed
+        with fabtools.files.watch('watched') as f:
+            pass
+        assert not f.changed
+
         # Ensure that the callable is triggered only
         # when the watched file is modified
         require.file('watched', contents='aaa')
-        with fabtools.files.watch('watched', False, require.file, 'modified1'):
+        with fabtools.files.watch('watched', callback=partial(require.file, 'modified1')):
             require.file('watched', contents='bbb')
         assert fabtools.files.is_file('modified1')
-        with fabtools.files.watch('watched', False, require.file, 'modified2'):
+        with fabtools.files.watch('watched', callback=partial(require.file, 'modified2')):
             pass
         assert not fabtools.files.is_file('modified2')