Simplify context managers for enabling units.
Fix doc warning
Merge pull request #1662 from mdboom/unit/simpler-context-manager

Simplify context manager for enabling units