From 37501d9ac3c451de95ec16aeec7678aaf23fa907 Mon Sep 17 00:00:00 2001
From: Antoine Millet <antoine.millet@smartjog.com>
Date: Mon, 3 Oct 2011 18:10:15 +0200
Subject: [PATCH] Added create_rpc and create_tunnel shortcuts with label auto
 assignation

---
 sjrpc/core/exceptions.py    |  8 ++++++++
 sjrpc/core/rpcconnection.py | 25 +++++++++++++++++++++----
 2 files changed, 29 insertions(+), 4 deletions(-)

diff --git a/sjrpc/core/exceptions.py b/sjrpc/core/exceptions.py
index 629e09d..5b2aa53 100644
--- a/sjrpc/core/exceptions.py
+++ b/sjrpc/core/exceptions.py
@@ -23,3 +23,11 @@ class SocketRpcError(Exception):
     '''
     Exception used internally to raise a socket fault.
     '''
+
+
+class NoFreeLabelError(Exception):
+
+    '''
+    Exception raised when no more free labels are available for protocol
+    allocation.
+    '''
diff --git a/sjrpc/core/rpcconnection.py b/sjrpc/core/rpcconnection.py
index 0a831f6..17f2a11 100644
--- a/sjrpc/core/rpcconnection.py
+++ b/sjrpc/core/rpcconnection.py
@@ -11,9 +11,8 @@ import struct
 import socket
 import logging
 
-from sjrpc.core.protocols.rpc import RpcProtocol
-from sjrpc.core.protocols import Protocol
-from sjrpc.core.exceptions import RpcError
+from sjrpc.core.protocols import Protocol, RpcProtocol, TunnelProtocol
+from sjrpc.core.exceptions import RpcError, NoFreeLabelError
 
 import pyev
 
@@ -54,6 +53,7 @@ class RpcConnection(object):
 
     MESSAGE_HEADER = '!HL'
     MESSAGE_HEADER_FALLBACK = '!L'
+    MAX_LABEL = 2 ** 16
     SHORTCUTS_MAINRPC = ('call', 'async_call')
 
     def __init__(self, sock, loop=None, enable_tcp_keepalive=False,
@@ -294,7 +294,12 @@ class RpcConnection(object):
         '''
         Register a new protocol for the specified label.
         '''
-
+        if label is None:
+            for label in xrange(0, RpcConnection.MAX_LABEL):
+                if label not in self._protocols:
+                    break
+            else:
+                raise NoFreeLabelError('No more label number are availables')
         if label in self._protocols:
             raise KeyError('A protocol is already registered for this label')
         elif not isinstance(label, int):
@@ -312,6 +317,18 @@ class RpcConnection(object):
         else:
             raise KeyError('No protocol registered for this label')
 
+    def create_rpc(self, label=None, *args, **kwargs):
+        '''
+        Shortcut which can be used to create rpc protocols.
+        '''
+        return self.register_protocol(label, RpcProtocol, *args, **kwargs)
+
+    def create_tunnel(self, label=None, *args, **kwargs):
+        '''
+        Shortcut which can be used to create tunnels protocols.
+        '''
+        return self.register_protocol(label, TunnelProtocol, *args, **kwargs)
+
     def get_protocol(self, label):
         '''
         Get the protocol registered for specified label.
-- 
GitLab