From f52bab51efffc94fd1acecd2534f254e81ca473d Mon Sep 17 00:00:00 2001
From: maralla <maralla.ai@gmail.com>
Date: Thu, 14 Apr 2016 17:05:03 +0800
Subject: [PATCH] expose socket handle in cython transports

---
 tests/test_cytransport.py                  | 18 ++++++++++++++++++
 thriftpy/transport/buffered/cybuffered.pyx |  1 -
 thriftpy/transport/cybase.pxd              |  2 ++
 thriftpy/transport/cybase.pyx              |  6 ++++++
 thriftpy/transport/framed/cyframed.pyx     |  1 -
 thriftpy/transport/memory/cymemory.pyx     |  1 +
 6 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/tests/test_cytransport.py b/tests/test_cytransport.py
index 5724aac..b0f4e18 100644
--- a/tests/test_cytransport.py
+++ b/tests/test_cytransport.py
@@ -34,3 +34,21 @@ def test_buffered_read():
     t.flush()
 
     assert t.read(4) == b"ping"
+
+
+def test_transport_handle():
+    from thriftpy._compat import CYTHON
+    if not CYTHON:
+        return
+
+    from thriftpy.transport import TSocket
+    from thriftpy.transport.memory import TCyMemoryBuffer
+
+    s = TSocket()
+    s.set_handle('the sock')
+
+    assert TCyBufferedTransport(s).sock == 'the sock'
+    assert TCyFramedTransport(s).sock == 'the sock'
+    assert TCyMemoryBuffer().sock is None
+    assert TCyBufferedTransport(TCyFramedTransport(s)).sock == 'the sock'
+    assert TCyBufferedTransport(TCyMemoryBuffer()).sock is None
diff --git a/thriftpy/transport/buffered/cybuffered.pyx b/thriftpy/transport/buffered/cybuffered.pyx
index 9364363..2264522 100644
--- a/thriftpy/transport/buffered/cybuffered.pyx
+++ b/thriftpy/transport/buffered/cybuffered.pyx
@@ -13,7 +13,6 @@ cdef class TCyBufferedTransport(CyTransportBase):
     """binary reader/writer"""
 
     cdef:
-        object trans
         TCyBuffer rbuf, wbuf
 
     def __init__(self, trans, int buf_size=DEFAULT_BUFFER):
diff --git a/thriftpy/transport/cybase.pxd b/thriftpy/transport/cybase.pxd
index f374638..a960be9 100644
--- a/thriftpy/transport/cybase.pxd
+++ b/thriftpy/transport/cybase.pxd
@@ -15,6 +15,8 @@ cdef class TCyBuffer(object):
 
 
 cdef class CyTransportBase(object):
+    cdef object trans
+
     cdef c_read(self, int sz, char* out)
     cdef c_write(self, char* data, int sz)
     cdef c_flush(self)
diff --git a/thriftpy/transport/cybase.pyx b/thriftpy/transport/cybase.pyx
index 8398c76..c05560f 100644
--- a/thriftpy/transport/cybase.pyx
+++ b/thriftpy/transport/cybase.pyx
@@ -115,6 +115,12 @@ cdef class CyTransportBase(object):
     def clean(self):
         pass
 
+    @property
+    def sock(self):
+        if not self.trans:
+            return
+        return getattr(self.trans, 'sock', None)
+
     cdef get_string(self, int sz):
         cdef:
             char out[STACK_STRING_LEN]
diff --git a/thriftpy/transport/framed/cyframed.pyx b/thriftpy/transport/framed/cyframed.pyx
index b43945c..250f4de 100644
--- a/thriftpy/transport/framed/cyframed.pyx
+++ b/thriftpy/transport/framed/cyframed.pyx
@@ -19,7 +19,6 @@ cdef extern from "../../protocol/cybin/endian_port.h":
 
 cdef class TCyFramedTransport(CyTransportBase):
     cdef:
-        object trans
         TCyBuffer rbuf, rframe_buf, wframe_buf
 
     def __init__(self, trans, int buf_size=DEFAULT_BUFFER):
diff --git a/thriftpy/transport/memory/cymemory.pyx b/thriftpy/transport/memory/cymemory.pyx
index 27ef05f..4270cd5 100644
--- a/thriftpy/transport/memory/cymemory.pyx
+++ b/thriftpy/transport/memory/cymemory.pyx
@@ -18,6 +18,7 @@ cdef class TCyMemoryBuffer(CyTransportBase):
     cdef TCyBuffer buf
 
     def __init__(self, value=b'', int buf_size=DEFAULT_BUFFER):
+        self.trans = None
         self.buf = TCyBuffer(buf_size)
 
         if value: