fixes

[ncclient] / ncclient / operations / lock.py

diff --git a/ncclient/operations/lock.py b/ncclient/operations/lock.py
index 533509e..d6cef80 100644 (file)
--- a/ncclient/operations/lock.py
+++ b/ncclient/operations/lock.py
@@ -14,11 +14,9 @@
 
 'Locking-related NETCONF operations'
 
-from copy import deepcopy
-
 from ncclient.rpc import RPC
 
-# TODO - a context manager around some <target> would be real neat
+import util
 
 class Lock(RPC):
     
@@ -31,7 +29,9 @@ class Lock(RPC):
     }
     
     def request(self, target='running'):
-        spec = deepcopy(Lock.SPEC)
+        if target=='candidate':
+            self._assert(':candidate')
+        spec = Lock.SPEC.copy()
         spec['children']['children']['tag'] = target
         return self._request(spec)
 
@@ -47,6 +47,23 @@ class Unlock(RPC):
     }
     
     def request(self, target='running'):
-        spec = deepcopy(Unlock.SPEC)
+        if target=='candidate':
+            self._assert(':candidate')
+        spec = Unlock.SPEC.copy()
         spec['children']['children']['tag'] = target
-        return self._request(self.spec)
+        return self._request(spec)
+
+
+class LockContext:
+        
+    def __init__(self, session, target='running'):
+        self.session = session
+        self.target = target
+        
+    def __enter__(self):
+        Lock(self.session).request(self.target)
+        return self
+    
+    def __exit__(self, t, v, tb):
+        Unlock(self.session).request(self.target)
+        return False