socket_options.py 4.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129
  1. # -*- coding: utf-8 -*-
  2. """The implementation of the SocketOptionsAdapter."""
  3. import socket
  4. import warnings
  5. import sys
  6. import requests
  7. from requests import adapters
  8. from .._compat import connection
  9. from .._compat import poolmanager
  10. from .. import exceptions as exc
  11. class SocketOptionsAdapter(adapters.HTTPAdapter):
  12. """An adapter for requests that allows users to specify socket options.
  13. Since version 2.4.0 of requests, it is possible to specify a custom list
  14. of socket options that need to be set before establishing the connection.
  15. Example usage::
  16. >>> import socket
  17. >>> import requests
  18. >>> from requests_toolbelt.adapters import socket_options
  19. >>> s = requests.Session()
  20. >>> opts = [(socket.IPPROTO_TCP, socket.TCP_NODELAY, 0)]
  21. >>> adapter = socket_options.SocketOptionsAdapter(socket_options=opts)
  22. >>> s.mount('http://', adapter)
  23. You can also take advantage of the list of default options on this class
  24. to keep using the original options in addition to your custom options. In
  25. that case, ``opts`` might look like::
  26. >>> opts = socket_options.SocketOptionsAdapter.default_options + opts
  27. """
  28. if connection is not None:
  29. default_options = getattr(
  30. connection.HTTPConnection,
  31. 'default_socket_options',
  32. [(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)]
  33. )
  34. else:
  35. default_options = []
  36. warnings.warn(exc.RequestsVersionTooOld,
  37. "This version of Requests is only compatible with a "
  38. "version of urllib3 which is too old to support "
  39. "setting options on a socket. This adapter is "
  40. "functionally useless.")
  41. def __init__(self, **kwargs):
  42. self.socket_options = kwargs.pop('socket_options',
  43. self.default_options)
  44. super(SocketOptionsAdapter, self).__init__(**kwargs)
  45. def init_poolmanager(self, connections, maxsize, block=False):
  46. if requests.__build__ >= 0x020400:
  47. # NOTE(Ian): Perhaps we should raise a warning
  48. self.poolmanager = poolmanager.PoolManager(
  49. num_pools=connections,
  50. maxsize=maxsize,
  51. block=block,
  52. socket_options=self.socket_options
  53. )
  54. else:
  55. super(SocketOptionsAdapter, self).init_poolmanager(
  56. connections, maxsize, block
  57. )
  58. class TCPKeepAliveAdapter(SocketOptionsAdapter):
  59. """An adapter for requests that turns on TCP Keep-Alive by default.
  60. The adapter sets 4 socket options:
  61. - ``SOL_SOCKET`` ``SO_KEEPALIVE`` - This turns on TCP Keep-Alive
  62. - ``IPPROTO_TCP`` ``TCP_KEEPINTVL`` 20 - Sets the keep alive interval
  63. - ``IPPROTO_TCP`` ``TCP_KEEPCNT`` 5 - Sets the number of keep alive probes
  64. - ``IPPROTO_TCP`` ``TCP_KEEPIDLE`` 60 - Sets the keep alive time if the
  65. socket library has the ``TCP_KEEPIDLE`` constant
  66. The latter three can be overridden by keyword arguments (respectively):
  67. - ``idle``
  68. - ``interval``
  69. - ``count``
  70. You can use this adapter like so::
  71. >>> from requests_toolbelt.adapters import socket_options
  72. >>> tcp = socket_options.TCPKeepAliveAdapter(idle=120, interval=10)
  73. >>> s = requests.Session()
  74. >>> s.mount('http://', tcp)
  75. """
  76. def __init__(self, **kwargs):
  77. socket_options = kwargs.pop('socket_options',
  78. SocketOptionsAdapter.default_options)
  79. idle = kwargs.pop('idle', 60)
  80. interval = kwargs.pop('interval', 20)
  81. count = kwargs.pop('count', 5)
  82. socket_options = socket_options + [
  83. (socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
  84. ]
  85. # NOTE(Ian): OSX does not have these constants defined, so we
  86. # set them conditionally.
  87. if getattr(socket, 'TCP_KEEPINTVL', None) is not None:
  88. socket_options += [(socket.IPPROTO_TCP, socket.TCP_KEEPINTVL,
  89. interval)]
  90. elif sys.platform == 'darwin':
  91. # On OSX, TCP_KEEPALIVE from netinet/tcp.h is not exported
  92. # by python's socket module
  93. TCP_KEEPALIVE = getattr(socket, 'TCP_KEEPALIVE', 0x10)
  94. socket_options += [(socket.IPPROTO_TCP, TCP_KEEPALIVE, interval)]
  95. if getattr(socket, 'TCP_KEEPCNT', None) is not None:
  96. socket_options += [(socket.IPPROTO_TCP, socket.TCP_KEEPCNT, count)]
  97. if getattr(socket, 'TCP_KEEPIDLE', None) is not None:
  98. socket_options += [(socket.IPPROTO_TCP, socket.TCP_KEEPIDLE, idle)]
  99. super(TCPKeepAliveAdapter, self).__init__(
  100. socket_options=socket_options, **kwargs
  101. )