PageRenderTime 52ms CodeModel.GetById 21ms RepoModel.GetById 0ms app.codeStats 1ms

/zmq/sugar/socket.py

http://github.com/zeromq/pyzmq
Python | 758 lines | 398 code | 49 blank | 311 comment | 52 complexity | 859df2a0285f84e324c8c1fc2582de04 MD5 | raw file
Possible License(s): BSD-3-Clause, LGPL-3.0, Apache-2.0
  1. # coding: utf-8
  2. """0MQ Socket pure Python methods."""
  3. # Copyright (C) PyZMQ Developers
  4. # Distributed under the terms of the Modified BSD License.
  5. import errno
  6. import random
  7. import sys
  8. import warnings
  9. import zmq
  10. from zmq.backend import Socket as SocketBase
  11. from .poll import Poller
  12. from . import constants
  13. from .attrsettr import AttributeSetter
  14. from zmq.error import ZMQError, ZMQBindError
  15. from zmq.utils import jsonapi
  16. from zmq.utils.strtypes import bytes, unicode, basestring
  17. from .constants import (
  18. SNDMORE, ENOTSUP, POLLIN,
  19. int64_sockopt_names,
  20. int_sockopt_names,
  21. bytes_sockopt_names,
  22. fd_sockopt_names,
  23. )
  24. try:
  25. import cPickle
  26. pickle = cPickle
  27. except:
  28. cPickle = None
  29. import pickle
  30. try:
  31. DEFAULT_PROTOCOL = pickle.DEFAULT_PROTOCOL
  32. except AttributeError:
  33. DEFAULT_PROTOCOL = pickle.HIGHEST_PROTOCOL
  34. class Socket(SocketBase, AttributeSetter):
  35. """The ZMQ socket object
  36. To create a Socket, first create a Context::
  37. ctx = zmq.Context.instance()
  38. then call ``ctx.socket(socket_type)``::
  39. s = ctx.socket(zmq.ROUTER)
  40. """
  41. _shadow = False
  42. _monitor_socket = None
  43. def __init__(self, *a, **kw):
  44. super(Socket, self).__init__(*a, **kw)
  45. if 'shadow' in kw:
  46. self._shadow = True
  47. else:
  48. self._shadow = False
  49. def __del__(self):
  50. if not self._shadow:
  51. self.close()
  52. # socket as context manager:
  53. def __enter__(self):
  54. """Sockets are context managers
  55. .. versionadded:: 14.4
  56. """
  57. return self
  58. def __exit__(self, *args, **kwargs):
  59. self.close()
  60. #-------------------------------------------------------------------------
  61. # Socket creation
  62. #-------------------------------------------------------------------------
  63. def __copy__(self, memo=None):
  64. """Copying a Socket creates a shadow copy"""
  65. return self.__class__.shadow(self.underlying)
  66. __deepcopy__ = __copy__
  67. @classmethod
  68. def shadow(cls, address):
  69. """Shadow an existing libzmq socket
  70. address is the integer address of the libzmq socket
  71. or an FFI pointer to it.
  72. .. versionadded:: 14.1
  73. """
  74. from zmq.utils.interop import cast_int_addr
  75. address = cast_int_addr(address)
  76. return cls(shadow=address)
  77. def close(self, linger=None):
  78. if self.context:
  79. self.context._rm_socket(self)
  80. super(Socket, self).close(linger=linger)
  81. #-------------------------------------------------------------------------
  82. # Deprecated aliases
  83. #-------------------------------------------------------------------------
  84. @property
  85. def socket_type(self):
  86. warnings.warn("Socket.socket_type is deprecated, use Socket.type",
  87. DeprecationWarning
  88. )
  89. return self.type
  90. #-------------------------------------------------------------------------
  91. # Hooks for sockopt completion
  92. #-------------------------------------------------------------------------
  93. def __dir__(self):
  94. keys = dir(self.__class__)
  95. for collection in (
  96. bytes_sockopt_names,
  97. int_sockopt_names,
  98. int64_sockopt_names,
  99. fd_sockopt_names,
  100. ):
  101. keys.extend(collection)
  102. return keys
  103. #-------------------------------------------------------------------------
  104. # Getting/Setting options
  105. #-------------------------------------------------------------------------
  106. setsockopt = SocketBase.set
  107. getsockopt = SocketBase.get
  108. def __setattr__(self, key, value):
  109. """Override to allow setting zmq.[UN]SUBSCRIBE even though we have a subscribe method"""
  110. if key in self.__dict__:
  111. object.__setattr__(self, key, value)
  112. return
  113. _key = key.lower()
  114. if _key in ('subscribe', 'unsubscribe'):
  115. if isinstance(value, unicode):
  116. value = value.encode('utf8')
  117. if _key == 'subscribe':
  118. self.set(zmq.SUBSCRIBE, value)
  119. else:
  120. self.set(zmq.UNSUBSCRIBE, value)
  121. return
  122. super(Socket, self).__setattr__(key, value)
  123. def fileno(self):
  124. """Return edge-triggered file descriptor for this socket.
  125. This is a read-only edge-triggered file descriptor for both read and write events on this socket.
  126. It is important that all available events be consumed when an event is detected,
  127. otherwise the read event will not trigger again.
  128. .. versionadded:: 17.0
  129. """
  130. return self.FD
  131. def subscribe(self, topic):
  132. """Subscribe to a topic
  133. Only for SUB sockets.
  134. .. versionadded:: 15.3
  135. """
  136. if isinstance(topic, unicode):
  137. topic = topic.encode('utf8')
  138. self.set(zmq.SUBSCRIBE, topic)
  139. def unsubscribe(self, topic):
  140. """Unsubscribe from a topic
  141. Only for SUB sockets.
  142. .. versionadded:: 15.3
  143. """
  144. if isinstance(topic, unicode):
  145. topic = topic.encode('utf8')
  146. self.set(zmq.UNSUBSCRIBE, topic)
  147. def set_string(self, option, optval, encoding='utf-8'):
  148. """Set socket options with a unicode object.
  149. This is simply a wrapper for setsockopt to protect from encoding ambiguity.
  150. See the 0MQ documentation for details on specific options.
  151. Parameters
  152. ----------
  153. option : int
  154. The name of the option to set. Can be any of: SUBSCRIBE,
  155. UNSUBSCRIBE, IDENTITY
  156. optval : unicode string (unicode on py2, str on py3)
  157. The value of the option to set.
  158. encoding : str
  159. The encoding to be used, default is utf8
  160. """
  161. if not isinstance(optval, unicode):
  162. raise TypeError("unicode strings only")
  163. return self.set(option, optval.encode(encoding))
  164. setsockopt_unicode = setsockopt_string = set_string
  165. def get_string(self, option, encoding='utf-8'):
  166. """Get the value of a socket option.
  167. See the 0MQ documentation for details on specific options.
  168. Parameters
  169. ----------
  170. option : int
  171. The option to retrieve.
  172. Returns
  173. -------
  174. optval : unicode string (unicode on py2, str on py3)
  175. The value of the option as a unicode string.
  176. """
  177. if option not in constants.bytes_sockopts:
  178. raise TypeError("option %i will not return a string to be decoded"%option)
  179. return self.getsockopt(option).decode(encoding)
  180. getsockopt_unicode = getsockopt_string = get_string
  181. def bind_to_random_port(self, addr, min_port=49152, max_port=65536, max_tries=100):
  182. """Bind this socket to a random port in a range.
  183. If the port range is unspecified, the system will choose the port.
  184. Parameters
  185. ----------
  186. addr : str
  187. The address string without the port to pass to ``Socket.bind()``.
  188. min_port : int, optional
  189. The minimum port in the range of ports to try (inclusive).
  190. max_port : int, optional
  191. The maximum port in the range of ports to try (exclusive).
  192. max_tries : int, optional
  193. The maximum number of bind attempts to make.
  194. Returns
  195. -------
  196. port : int
  197. The port the socket was bound to.
  198. Raises
  199. ------
  200. ZMQBindError
  201. if `max_tries` reached before successful bind
  202. """
  203. if hasattr(constants, 'LAST_ENDPOINT') and min_port == 49152 and max_port == 65536:
  204. # if LAST_ENDPOINT is supported, and min_port / max_port weren't specified,
  205. # we can bind to port 0 and let the OS do the work
  206. self.bind("%s:*" % addr)
  207. url = self.last_endpoint.decode('ascii', 'replace')
  208. _, port_s = url.rsplit(':', 1)
  209. return int(port_s)
  210. for i in range(max_tries):
  211. try:
  212. port = random.randrange(min_port, max_port)
  213. self.bind('%s:%s' % (addr, port))
  214. except ZMQError as exception:
  215. en = exception.errno
  216. if en == zmq.EADDRINUSE:
  217. continue
  218. elif sys.platform == 'win32' and en == errno.EACCES:
  219. continue
  220. else:
  221. raise
  222. else:
  223. return port
  224. raise ZMQBindError("Could not bind socket to random port.")
  225. def get_hwm(self):
  226. """Get the High Water Mark.
  227. On libzmq ≥ 3, this gets SNDHWM if available, otherwise RCVHWM
  228. """
  229. major = zmq.zmq_version_info()[0]
  230. if major >= 3:
  231. # return sndhwm, fallback on rcvhwm
  232. try:
  233. return self.getsockopt(zmq.SNDHWM)
  234. except zmq.ZMQError:
  235. pass
  236. return self.getsockopt(zmq.RCVHWM)
  237. else:
  238. return self.getsockopt(zmq.HWM)
  239. def set_hwm(self, value):
  240. """Set the High Water Mark.
  241. On libzmq ≥ 3, this sets both SNDHWM and RCVHWM
  242. .. warning::
  243. New values only take effect for subsequent socket
  244. bind/connects.
  245. """
  246. major = zmq.zmq_version_info()[0]
  247. if major >= 3:
  248. raised = None
  249. try:
  250. self.sndhwm = value
  251. except Exception as e:
  252. raised = e
  253. try:
  254. self.rcvhwm = value
  255. except Exception as e:
  256. raised = e
  257. if raised:
  258. raise raised
  259. else:
  260. return self.setsockopt(zmq.HWM, value)
  261. hwm = property(get_hwm, set_hwm,
  262. """Property for High Water Mark.
  263. Setting hwm sets both SNDHWM and RCVHWM as appropriate.
  264. It gets SNDHWM if available, otherwise RCVHWM.
  265. """
  266. )
  267. #-------------------------------------------------------------------------
  268. # Sending and receiving messages
  269. #-------------------------------------------------------------------------
  270. def send(self, data, flags=0, copy=True, track=False, routing_id=None, group=None):
  271. """Send a single zmq message frame on this socket.
  272. This queues the message to be sent by the IO thread at a later time.
  273. With flags=NOBLOCK, this raises :class:`ZMQError` if the queue is full;
  274. otherwise, this waits until space is available.
  275. See :class:`Poller` for more general non-blocking I/O.
  276. Parameters
  277. ----------
  278. data : bytes, Frame, memoryview
  279. The content of the message. This can be any object that provides
  280. the Python buffer API (i.e. `memoryview(data)` can be called).
  281. flags : int
  282. 0, NOBLOCK, SNDMORE, or NOBLOCK|SNDMORE.
  283. copy : bool
  284. Should the message be sent in a copying or non-copying manner.
  285. track : bool
  286. Should the message be tracked for notification that ZMQ has
  287. finished with it? (ignored if copy=True)
  288. routing_id : int
  289. For use with SERVER sockets
  290. group : str
  291. For use with RADIO sockets
  292. Returns
  293. -------
  294. None : if `copy` or not track
  295. None if message was sent, raises an exception otherwise.
  296. MessageTracker : if track and not copy
  297. a MessageTracker object, whose `pending` property will
  298. be True until the send is completed.
  299. Raises
  300. ------
  301. TypeError
  302. If a unicode object is passed
  303. ValueError
  304. If `track=True`, but an untracked Frame is passed.
  305. ZMQError
  306. If the send does not succeed for any reason (including
  307. if NOBLOCK is set and the outgoing queue is full).
  308. .. versionchanged:: 17.0
  309. DRAFT support for routing_id and group arguments.
  310. """
  311. if routing_id is not None:
  312. if not isinstance(data, zmq.Frame):
  313. data = zmq.Frame(data, track=track, copy=copy or None,
  314. copy_threshold=self.copy_threshold)
  315. data.routing_id = routing_id
  316. if group is not None:
  317. if not isinstance(data, zmq.Frame):
  318. data = zmq.Frame(data, track=track, copy=copy or None,
  319. copy_threshold=self.copy_threshold)
  320. data.group = group
  321. return super(Socket, self).send(data, flags=flags, copy=copy, track=track)
  322. def send_multipart(self, msg_parts, flags=0, copy=True, track=False, **kwargs):
  323. """Send a sequence of buffers as a multipart message.
  324. The zmq.SNDMORE flag is added to all msg parts before the last.
  325. Parameters
  326. ----------
  327. msg_parts : iterable
  328. A sequence of objects to send as a multipart message. Each element
  329. can be any sendable object (Frame, bytes, buffer-providers)
  330. flags : int, optional
  331. Any valid flags for :func:`Socket.send`.
  332. SNDMORE is added automatically for frames before the last.
  333. copy : bool, optional
  334. Should the frame(s) be sent in a copying or non-copying manner.
  335. If copy=False, frames smaller than self.copy_threshold bytes
  336. will be copied anyway.
  337. track : bool, optional
  338. Should the frame(s) be tracked for notification that ZMQ has
  339. finished with it (ignored if copy=True).
  340. Returns
  341. -------
  342. None : if copy or not track
  343. MessageTracker : if track and not copy
  344. a MessageTracker object, whose `pending` property will
  345. be True until the last send is completed.
  346. """
  347. # typecheck parts before sending:
  348. for i,msg in enumerate(msg_parts):
  349. if isinstance(msg, (zmq.Frame, bytes, memoryview)):
  350. continue
  351. try:
  352. memoryview(msg)
  353. except Exception:
  354. rmsg = repr(msg)
  355. if len(rmsg) > 32:
  356. rmsg = rmsg[:32] + '...'
  357. raise TypeError(
  358. "Frame %i (%s) does not support the buffer interface." % (
  359. i, rmsg,
  360. ))
  361. for msg in msg_parts[:-1]:
  362. self.send(msg, SNDMORE|flags, copy=copy, track=track)
  363. # Send the last part without the extra SNDMORE flag.
  364. return self.send(msg_parts[-1], flags, copy=copy, track=track)
  365. def recv_multipart(self, flags=0, copy=True, track=False):
  366. """Receive a multipart message as a list of bytes or Frame objects
  367. Parameters
  368. ----------
  369. flags : int, optional
  370. Any valid flags for :func:`Socket.recv`.
  371. copy : bool, optional
  372. Should the message frame(s) be received in a copying or non-copying manner?
  373. If False a Frame object is returned for each part, if True a copy of
  374. the bytes is made for each frame.
  375. track : bool, optional
  376. Should the message frame(s) be tracked for notification that ZMQ has
  377. finished with it? (ignored if copy=True)
  378. Returns
  379. -------
  380. msg_parts : list
  381. A list of frames in the multipart message; either Frames or bytes,
  382. depending on `copy`.
  383. Raises
  384. ------
  385. ZMQError
  386. for any of the reasons :func:`~Socket.recv` might fail
  387. """
  388. parts = [self.recv(flags, copy=copy, track=track)]
  389. # have first part already, only loop while more to receive
  390. while self.getsockopt(zmq.RCVMORE):
  391. part = self.recv(flags, copy=copy, track=track)
  392. parts.append(part)
  393. return parts
  394. def _deserialize(self, recvd, load):
  395. """Deserialize a received message
  396. Override in subclass (e.g. Futures) if recvd is not the raw bytes.
  397. The default implementation expects bytes and returns the deserialized message immediately.
  398. Parameters
  399. ----------
  400. load: callable
  401. Callable that deserializes bytes
  402. recvd:
  403. The object returned by self.recv
  404. """
  405. return load(recvd)
  406. def send_serialized(self, msg, serialize, flags=0, copy=True, **kwargs):
  407. """Send a message with a custom serialization function.
  408. .. versionadded:: 17
  409. Parameters
  410. ----------
  411. msg : The message to be sent. Can be any object serializable by `serialize`.
  412. serialize : callable
  413. The serialization function to use.
  414. serialize(msg) should return an iterable of sendable message frames
  415. (e.g. bytes objects), which will be passed to send_multipart.
  416. flags : int, optional
  417. Any valid flags for :func:`Socket.send`.
  418. copy : bool, optional
  419. Whether to copy the frames.
  420. """
  421. frames = serialize(msg)
  422. return self.send_multipart(frames, flags=flags, copy=copy, **kwargs)
  423. def recv_serialized(self, deserialize, flags=0, copy=True):
  424. """Receive a message with a custom deserialization function.
  425. .. versionadded:: 17
  426. Parameters
  427. ----------
  428. deserialize : callable
  429. The deserialization function to use.
  430. deserialize will be called with one argument: the list of frames
  431. returned by recv_multipart() and can return any object.
  432. flags : int, optional
  433. Any valid flags for :func:`Socket.recv`.
  434. copy : bool, optional
  435. Whether to recv bytes or Frame objects.
  436. Returns
  437. -------
  438. obj : object
  439. The object returned by the deserialization function.
  440. Raises
  441. ------
  442. ZMQError
  443. for any of the reasons :func:`~Socket.recv` might fail
  444. """
  445. frames = self.recv_multipart(flags=flags, copy=copy)
  446. return self._deserialize(frames, deserialize)
  447. def send_string(self, u, flags=0, copy=True, encoding='utf-8', **kwargs):
  448. """Send a Python unicode string as a message with an encoding.
  449. 0MQ communicates with raw bytes, so you must encode/decode
  450. text (unicode on py2, str on py3) around 0MQ.
  451. Parameters
  452. ----------
  453. u : Python unicode string (unicode on py2, str on py3)
  454. The unicode string to send.
  455. flags : int, optional
  456. Any valid flags for :func:`Socket.send`.
  457. encoding : str [default: 'utf-8']
  458. The encoding to be used
  459. """
  460. if not isinstance(u, basestring):
  461. raise TypeError("unicode/str objects only")
  462. return self.send(u.encode(encoding), flags=flags, copy=copy, **kwargs)
  463. send_unicode = send_string
  464. def recv_string(self, flags=0, encoding='utf-8'):
  465. """Receive a unicode string, as sent by send_string.
  466. Parameters
  467. ----------
  468. flags : int
  469. Any valid flags for :func:`Socket.recv`.
  470. encoding : str [default: 'utf-8']
  471. The encoding to be used
  472. Returns
  473. -------
  474. s : unicode string (unicode on py2, str on py3)
  475. The Python unicode string that arrives as encoded bytes.
  476. Raises
  477. ------
  478. ZMQError
  479. for any of the reasons :func:`~Socket.recv` might fail
  480. """
  481. msg = self.recv(flags=flags)
  482. return self._deserialize(msg, lambda buf: buf.decode(encoding))
  483. recv_unicode = recv_string
  484. def send_pyobj(self, obj, flags=0, protocol=DEFAULT_PROTOCOL, **kwargs):
  485. """Send a Python object as a message using pickle to serialize.
  486. Parameters
  487. ----------
  488. obj : Python object
  489. The Python object to send.
  490. flags : int
  491. Any valid flags for :func:`Socket.send`.
  492. protocol : int
  493. The pickle protocol number to use. The default is pickle.DEFAULT_PROTOCOL
  494. where defined, and pickle.HIGHEST_PROTOCOL elsewhere.
  495. """
  496. msg = pickle.dumps(obj, protocol)
  497. return self.send(msg, flags=flags, **kwargs)
  498. def recv_pyobj(self, flags=0):
  499. """Receive a Python object as a message using pickle to serialize.
  500. Parameters
  501. ----------
  502. flags : int
  503. Any valid flags for :func:`Socket.recv`.
  504. Returns
  505. -------
  506. obj : Python object
  507. The Python object that arrives as a message.
  508. Raises
  509. ------
  510. ZMQError
  511. for any of the reasons :func:`~Socket.recv` might fail
  512. """
  513. msg = self.recv(flags)
  514. return self._deserialize(msg, pickle.loads)
  515. def send_json(self, obj, flags=0, **kwargs):
  516. """Send a Python object as a message using json to serialize.
  517. Keyword arguments are passed on to json.dumps
  518. Parameters
  519. ----------
  520. obj : Python object
  521. The Python object to send
  522. flags : int
  523. Any valid flags for :func:`Socket.send`
  524. """
  525. send_kwargs = {}
  526. for key in ('routing_id', 'group'):
  527. if key in kwargs:
  528. send_kwargs[key] = kwargs.pop(key)
  529. msg = jsonapi.dumps(obj, **kwargs)
  530. return self.send(msg, flags=flags, **send_kwargs)
  531. def recv_json(self, flags=0, **kwargs):
  532. """Receive a Python object as a message using json to serialize.
  533. Keyword arguments are passed on to json.loads
  534. Parameters
  535. ----------
  536. flags : int
  537. Any valid flags for :func:`Socket.recv`.
  538. Returns
  539. -------
  540. obj : Python object
  541. The Python object that arrives as a message.
  542. Raises
  543. ------
  544. ZMQError
  545. for any of the reasons :func:`~Socket.recv` might fail
  546. """
  547. msg = self.recv(flags)
  548. return self._deserialize(msg, lambda buf: jsonapi.loads(buf, **kwargs))
  549. _poller_class = Poller
  550. def poll(self, timeout=None, flags=POLLIN):
  551. """Poll the socket for events.
  552. See :class:`Poller` to wait for multiple sockets at once.
  553. Parameters
  554. ----------
  555. timeout : int [default: None]
  556. The timeout (in milliseconds) to wait for an event. If unspecified
  557. (or specified None), will wait forever for an event.
  558. flags : int [default: POLLIN]
  559. POLLIN, POLLOUT, or POLLIN|POLLOUT. The event flags to poll for.
  560. Returns
  561. -------
  562. events : int
  563. The events that are ready and waiting,
  564. 0 if the timeout was reached with no events.
  565. """
  566. if self.closed:
  567. raise ZMQError(ENOTSUP)
  568. p = self._poller_class()
  569. p.register(self, flags)
  570. evts = dict(p.poll(timeout))
  571. # return 0 if no events, otherwise return event bitfield
  572. return evts.get(self, 0)
  573. def get_monitor_socket(self, events=None, addr=None):
  574. """Return a connected PAIR socket ready to receive the event notifications.
  575. .. versionadded:: libzmq-4.0
  576. .. versionadded:: 14.0
  577. Parameters
  578. ----------
  579. events : int [default: ZMQ_EVENT_ALL]
  580. The bitmask defining which events are wanted.
  581. addr : string [default: None]
  582. The optional endpoint for the monitoring sockets.
  583. Returns
  584. -------
  585. socket : (PAIR)
  586. The socket is already connected and ready to receive messages.
  587. """
  588. # safe-guard, method only available on libzmq >= 4
  589. if zmq.zmq_version_info() < (4,):
  590. raise NotImplementedError("get_monitor_socket requires libzmq >= 4, have %s" % zmq.zmq_version())
  591. # if already monitoring, return existing socket
  592. if self._monitor_socket:
  593. if self._monitor_socket.closed:
  594. self._monitor_socket = None
  595. else:
  596. return self._monitor_socket
  597. if addr is None:
  598. # create endpoint name from internal fd
  599. addr = "inproc://monitor.s-%d" % self.FD
  600. if events is None:
  601. # use all events
  602. events = zmq.EVENT_ALL
  603. # attach monitoring socket
  604. self.monitor(addr, events)
  605. # create new PAIR socket and connect it
  606. self._monitor_socket = self.context.socket(zmq.PAIR)
  607. self._monitor_socket.connect(addr)
  608. return self._monitor_socket
  609. def disable_monitor(self):
  610. """Shutdown the PAIR socket (created using get_monitor_socket)
  611. that is serving socket events.
  612. .. versionadded:: 14.4
  613. """
  614. self._monitor_socket = None
  615. self.monitor(None, 0)
  616. __all__ = ['Socket']