Simplify context manager for enabling units

It has always bugged me that we had separate functions for enabling
units permanently and temporarily.  This merges `set_enabled_units`
and `set_enabled_units_context` into a single function by using a
custom context manager rather than `contextlib.contextmanager`.

It would be great to have this small change in 0.3, as it affects API.