fixes

[ncclient] / ncclient / operations / lock.py

diff --git a/ncclient/operations/lock.py b/ncclient/operations/lock.py
index ed818a5..d6cef80 100644 (file)
--- a/ncclient/operations/lock.py
+++ b/ncclient/operations/lock.py
@@ -14,32 +14,56 @@
 
 'Locking-related NETCONF operations'
 
-from rpc import RPC
+from ncclient.rpc import RPC
 
+import util
 
 class Lock(RPC):
     
-    def __init__(self, session):
-        RPC.__init__(self, session)
-        self.spec = {
-            'tag': 'lock',
-            'children': { 'tag': 'target', 'children': {'tag': None} }
-            }
+    SPEC = {
+        'tag': 'lock',
+        'children': {
+            'tag': 'target',
+            'children': {'tag': None }
+        }
+    }
     
-    def request(self, target='running', reply_event=None):
-        self.spec['children']['children']['tag'] = target
-        self._do_request(self.spec, reply_event)
+    def request(self, target='running'):
+        if target=='candidate':
+            self._assert(':candidate')
+        spec = Lock.SPEC.copy()
+        spec['children']['children']['tag'] = target
+        return self._request(spec)
 
 
 class Unlock(RPC):
     
-    def __init__(self, session):
-        RPC.__init__(self, session)
-        self.spec = {
-            'tag': 'unlock',
-            'children': { 'tag': 'target', 'children': {'tag': None} }
-            }
+    SPEC = {
+        'tag': 'unlock',
+        'children': {
+            'tag': 'target',
+            'children': {'tag': None }
+        }
+    }
     
-    def request(self, target='running', reply_event=None):
-        self.spec['children']['children']['tag'] = target
-        self._do_request(self.spec, reply_event)
+    def request(self, target='running'):
+        if target=='candidate':
+            self._assert(':candidate')
+        spec = Unlock.SPEC.copy()
+        spec['children']['children']['tag'] = target
+        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