no need for weakvaluedict in RPCReplyListener, after all

[ncclient] / ncclient / operations / lock.py

diff --git a/ncclient/operations/lock.py b/ncclient/operations/lock.py
index 9af41f6..c6cb848 100644 (file)
--- a/ncclient/operations/lock.py
+++ b/ncclient/operations/lock.py
@@ -1,4 +1,4 @@
-# Copyright 2009 Shikhar Bhushan
+# Copyright 2h009 Shikhar Bhushan
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -14,31 +14,103 @@
 
 'Locking-related NETCONF operations'
 
-from rpc import RPC
+from copy import deepcopy
 
+from rpc import RPC, RPCReply, RPCError
+
+#class LockReply(RPCReply):
+#
+#    ERROR_CLS = LockDeniedError
+#
+#class LockDeniedError(RPCError):
+#
+#    def __new__(cls, err_dict):
+#        if rpcerr['tag'] != 'lock-denied':
+#            return RPCError(err_dict)
+#        else:
+#            return object.__new__(LockDeniedError)
+#
+#    def __init__(self, err_dict):
+#        RPCError.__init__(self, err_dict)
 
 class Lock(RPC):
-    
-    def __init__(self):
-        RPC.__init__(self)
-        self.spec = {
-            'tag': 'lock',
-            'children': [ { 'tag': 'target', 'text': None } ]
-            }
-    
-    def request(self, target='running', reply_event=None):
-        self.spec['children'][0]['text'] = target
-        self._do_request(self.spec, reply_event)
-
-class Unock(RPC):
-    
-    def __init__(self):
-        RPC.__init__(self)
-        self.spec = {
-            'tag': 'unlock',
-            'children': [ { 'tag': 'target', 'text': None } ]
-            }
-    
-    def request(self, target='running', reply_event=None):
-        self.spec['children'][0]['text'] = target
-        self._do_request(self.spec, reply_event)
+
+    # TESTED
+
+    "*<lock>* RPC"
+
+    SPEC = {
+        'tag': 'lock',
+        'subtree': {
+            'tag': 'target',
+            'subtree': {'tag': None }
+        }
+    }
+
+    #REPLY_CLS = LockReply
+
+    def request(self, target):
+        """
+        :arg target: see :ref:`source_target`
+        :type target: string
+
+        :rtype: :ref:`return`
+        """
+        spec = deepcopy(Lock.SPEC)
+        spec['subtree']['subtree']['tag'] = target
+        return self._request(spec)
+
+
+class Unlock(RPC):
+
+    # TESTED
+
+    "*<unlock>* RPC"
+
+    SPEC = {
+        'tag': 'unlock',
+        'subtree': {
+            'tag': 'target',
+            'subtree': {'tag': None }
+        }
+    }
+
+    def request(self, target):
+        """
+        :arg target: see :ref:`source_target`
+        :type target: string
+
+        :rtype: :ref:`return`
+        """
+        spec = deepcopy(Unlock.SPEC)
+        spec['subtree']['subtree']['tag'] = target
+        return self._request(spec)
+
+
+class LockContext:
+
+    # TESTED
+
+    """
+    A context manager for the :class:`Lock` / :class:`Unlock` pair of RPC's.
+
+    Initialise with session instance (:class:`Session
+    <ncclient.transport.Session>`) and lock target (:ref:`source_target`)
+    """
+
+    def __init__(self, session, target):
+        self.session = session
+        self.target = target
+
+    def __enter__(self):
+        reply = Lock(self.session).request(self.target)
+        if not reply.ok:
+            raise reply.error
+        else:
+            return self
+
+    def __exit__(self, *args):
+        reply = Unlock(self.session).request(self.target)
+        if not reply.ok:
+            raise reply.error
+        return False