catacomb/pwsafe.py: Make `PW' be a context manager, and use it.

[catacomb-python] / catacomb / pwsafe.py

diff --git a/catacomb/pwsafe.py b/catacomb/pwsafe.py
index 1e8e4dbfeb41ce1097be5868c2a1b985922a658d..82f2a46f1df6952c202e18974b471cc8130b5408 100644 (file)
--- a/catacomb/pwsafe.py
+++ b/catacomb/pwsafe.py
@@ -299,4 +299,11 @@ class PW (object):
       if k[0] == '$': yield me.unpack(me.db[k])[0]
       k = me.db.nextkey(k)
 
       if k[0] == '$': yield me.unpack(me.db[k])[0]
       k = me.db.nextkey(k)
 

+  ## Context protocol.
+
+  def __enter__(me):
+    return me
+  def __exit__(me, excty, excval, exctb):
+    me.db.close()
+

 ###----- That's all, folks --------------------------------------------------
 ###----- That's all, folks --------------------------------------------------