Format: text/x-python
⬇ Download File (26.4 KB)
   1  """Wrappers for forwarding stdout/stderr over zmq"""
   2  
   3  # Copyright (c) IPython Development Team.
   4  # Distributed under the terms of the Modified BSD License.
   5  
   6  import asyncio
   7  import atexit
   8  import contextvars
   9  import io
  10  import os
  11  import sys
  12  import threading
  13  import traceback
  14  import warnings
  15  from binascii import b2a_hex
  16  from collections import defaultdict, deque
  17  from io import StringIO, TextIOBase
  18  from threading import local
  19  from typing import Any, Callable, Deque, Dict, Optional
  20  
  21  import zmq
  22  from jupyter_client.session import extract_header
  23  from tornado.ioloop import IOLoop
  24  from zmq.eventloop.zmqstream import ZMQStream
  25  
  26  # -----------------------------------------------------------------------------
  27  # Globals
  28  # -----------------------------------------------------------------------------
  29  
  30  MASTER = 0
  31  CHILD = 1
  32  
  33  PIPE_BUFFER_SIZE = 1000
  34  
  35  # -----------------------------------------------------------------------------
  36  # IO classes
  37  # -----------------------------------------------------------------------------
  38  
  39  
  40  class IOPubThread:
  41      """An object for sending IOPub messages in a background thread
  42  
  43      Prevents a blocking main thread from delaying output from threads.
  44  
  45      IOPubThread(pub_socket).background_socket is a Socket-API-providing object
  46      whose IO is always run in a thread.
  47      """
  48  
  49      def __init__(self, socket, pipe=False):
  50          """Create IOPub thread
  51  
  52          Parameters
  53          ----------
  54          socket : zmq.PUB Socket
  55              the socket on which messages will be sent.
  56          pipe : bool
  57              Whether this process should listen for IOPub messages
  58              piped from subprocesses.
  59          """
  60          self.socket = socket
  61          self._stopped = False
  62          self.background_socket = BackgroundSocket(self)
  63          self._master_pid = os.getpid()
  64          self._pipe_flag = pipe
  65          self.io_loop = IOLoop(make_current=False)
  66          if pipe:
  67              self._setup_pipe_in()
  68          self._local = threading.local()
  69          self._events: Deque[Callable[..., Any]] = deque()
  70          self._event_pipes: Dict[threading.Thread, Any] = {}
  71          self._event_pipe_gc_lock: threading.Lock = threading.Lock()
  72          self._event_pipe_gc_seconds: float = 10
  73          self._event_pipe_gc_task: Optional[asyncio.Task[Any]] = None
  74          self._setup_event_pipe()
  75          self.thread = threading.Thread(target=self._thread_main, name="IOPub")
  76          self.thread.daemon = True
  77          self.thread.pydev_do_not_trace = True  # type:ignore[attr-defined]
  78          self.thread.is_pydev_daemon_thread = True  # type:ignore[attr-defined]
  79          self.thread.name = "IOPub"
  80  
  81      def _thread_main(self):
  82          """The inner loop that's actually run in a thread"""
  83  
  84          def _start_event_gc():
  85              self._event_pipe_gc_task = asyncio.ensure_future(self._run_event_pipe_gc())
  86  
  87          self.io_loop.run_sync(_start_event_gc)
  88  
  89          if not self._stopped:
  90              # avoid race if stop called before start thread gets here
  91              # probably only comes up in tests
  92              self.io_loop.start()
  93  
  94          if self._event_pipe_gc_task is not None:
  95              # cancel gc task to avoid pending task warnings
  96              async def _cancel():
  97                  self._event_pipe_gc_task.cancel()  # type:ignore[union-attr]
  98  
  99              if not self._stopped:
 100                  self.io_loop.run_sync(_cancel)
 101              else:
 102                  self._event_pipe_gc_task.cancel()
 103  
 104          self.io_loop.close(all_fds=True)
 105  
 106      def _setup_event_pipe(self):
 107          """Create the PULL socket listening for events that should fire in this thread."""
 108          ctx = self.socket.context
 109          pipe_in = ctx.socket(zmq.PULL)
 110          pipe_in.linger = 0
 111  
 112          _uuid = b2a_hex(os.urandom(16)).decode("ascii")
 113          iface = self._event_interface = "inproc://%s" % _uuid
 114          pipe_in.bind(iface)
 115          self._event_puller = ZMQStream(pipe_in, self.io_loop)
 116          self._event_puller.on_recv(self._handle_event)
 117  
 118      async def _run_event_pipe_gc(self):
 119          """Task to run event pipe gc continuously"""
 120          while True:
 121              await asyncio.sleep(self._event_pipe_gc_seconds)
 122              try:
 123                  await self._event_pipe_gc()
 124              except Exception as e:
 125                  print(f"Exception in IOPubThread._event_pipe_gc: {e}", file=sys.__stderr__)
 126  
 127      async def _event_pipe_gc(self):
 128          """run a single garbage collection on event pipes"""
 129          if not self._event_pipes:
 130              # don't acquire the lock if there's nothing to do
 131              return
 132          with self._event_pipe_gc_lock:
 133              for thread, socket in list(self._event_pipes.items()):
 134                  if not thread.is_alive():
 135                      socket.close()
 136                      del self._event_pipes[thread]
 137  
 138      @property
 139      def _event_pipe(self):
 140          """thread-local event pipe for signaling events that should be processed in the thread"""
 141          try:
 142              event_pipe = self._local.event_pipe
 143          except AttributeError:
 144              # new thread, new event pipe
 145              ctx = self.socket.context
 146              event_pipe = ctx.socket(zmq.PUSH)
 147              event_pipe.linger = 0
 148              event_pipe.connect(self._event_interface)
 149              self._local.event_pipe = event_pipe
 150              # associate event pipes to their threads
 151              # so they can be closed explicitly
 152              # implicit close on __del__ throws a ResourceWarning
 153              with self._event_pipe_gc_lock:
 154                  self._event_pipes[threading.current_thread()] = event_pipe
 155          return event_pipe
 156  
 157      def _handle_event(self, msg):
 158          """Handle an event on the event pipe
 159  
 160          Content of the message is ignored.
 161  
 162          Whenever *an* event arrives on the event stream,
 163          *all* waiting events are processed in order.
 164          """
 165          # freeze event count so new writes don't extend the queue
 166          # while we are processing
 167          n_events = len(self._events)
 168          for _ in range(n_events):
 169              event_f = self._events.popleft()
 170              event_f()
 171  
 172      def _setup_pipe_in(self):
 173          """setup listening pipe for IOPub from forked subprocesses"""
 174          ctx = self.socket.context
 175  
 176          # use UUID to authenticate pipe messages
 177          self._pipe_uuid = os.urandom(16)
 178  
 179          pipe_in = ctx.socket(zmq.PULL)
 180          pipe_in.linger = 0
 181  
 182          try:
 183              self._pipe_port = pipe_in.bind_to_random_port("tcp://127.0.0.1")
 184          except zmq.ZMQError as e:
 185              warnings.warn(
 186                  "Couldn't bind IOPub Pipe to 127.0.0.1: %s" % e
 187                  + "\nsubprocess output will be unavailable.",
 188                  stacklevel=2,
 189              )
 190              self._pipe_flag = False
 191              pipe_in.close()
 192              return
 193          self._pipe_in = ZMQStream(pipe_in, self.io_loop)
 194          self._pipe_in.on_recv(self._handle_pipe_msg)
 195  
 196      def _handle_pipe_msg(self, msg):
 197          """handle a pipe message from a subprocess"""
 198          if not self._pipe_flag or not self._is_master_process():
 199              return
 200          if msg[0] != self._pipe_uuid:
 201              print("Bad pipe message: %s", msg, file=sys.__stderr__)
 202              return
 203          self.send_multipart(msg[1:])
 204  
 205      def _setup_pipe_out(self):
 206          # must be new context after fork
 207          ctx = zmq.Context()
 208          pipe_out = ctx.socket(zmq.PUSH)
 209          pipe_out.linger = 3000  # 3s timeout for pipe_out sends before discarding the message
 210          pipe_out.connect("tcp://127.0.0.1:%i" % self._pipe_port)
 211          return ctx, pipe_out
 212  
 213      def _is_master_process(self):
 214          return os.getpid() == self._master_pid
 215  
 216      def _check_mp_mode(self):
 217          """check for forks, and switch to zmq pipeline if necessary"""
 218          if not self._pipe_flag or self._is_master_process():
 219              return MASTER
 220          return CHILD
 221  
 222      def start(self):
 223          """Start the IOPub thread"""
 224          self.thread.name = "IOPub"
 225          self.thread.start()
 226          # make sure we don't prevent process exit
 227          # I'm not sure why setting daemon=True above isn't enough, but it doesn't appear to be.
 228          atexit.register(self.stop)
 229  
 230      def stop(self):
 231          """Stop the IOPub thread"""
 232          self._stopped = True
 233          if not self.thread.is_alive():
 234              return
 235          self.io_loop.add_callback(self.io_loop.stop)
 236  
 237          self.thread.join(timeout=30)
 238          if self.thread.is_alive():
 239              # avoid infinite hang if stop fails
 240              msg = "IOPub thread did not terminate in 30 seconds"
 241              raise TimeoutError(msg)
 242          # close *all* event pipes, created in any thread
 243          # event pipes can only be used from other threads while self.thread.is_alive()
 244          # so after thread.join, this should be safe
 245          for _thread, event_pipe in self._event_pipes.items():
 246              event_pipe.close()
 247  
 248      def close(self):
 249          """Close the IOPub thread."""
 250          if self.closed:
 251              return
 252          self.socket.close()
 253          self.socket = None
 254  
 255      @property
 256      def closed(self):
 257          return self.socket is None
 258  
 259      def schedule(self, f):
 260          """Schedule a function to be called in our IO thread.
 261  
 262          If the thread is not running, call immediately.
 263          """
 264          if self.thread.is_alive():
 265              self._events.append(f)
 266              # wake event thread (message content is ignored)
 267              self._event_pipe.send(b"")
 268          else:
 269              f()
 270  
 271      def send_multipart(self, *args, **kwargs):
 272          """send_multipart schedules actual zmq send in my thread.
 273  
 274          If my thread isn't running (e.g. forked process), send immediately.
 275          """
 276          self.schedule(lambda: self._really_send(*args, **kwargs))
 277  
 278      def _really_send(self, msg, *args, **kwargs):
 279          """The callback that actually sends messages"""
 280          if self.closed:
 281              return
 282  
 283          mp_mode = self._check_mp_mode()
 284  
 285          if mp_mode != CHILD:
 286              # we are master, do a regular send
 287              self.socket.send_multipart(msg, *args, **kwargs)
 288          else:
 289              # we are a child, pipe to master
 290              # new context/socket for every pipe-out
 291              # since forks don't teardown politely, use ctx.term to ensure send has completed
 292              ctx, pipe_out = self._setup_pipe_out()
 293              pipe_out.send_multipart([self._pipe_uuid, *msg], *args, **kwargs)
 294              pipe_out.close()
 295              ctx.term()
 296  
 297  
 298  class BackgroundSocket:
 299      """Wrapper around IOPub thread that provides zmq send[_multipart]"""
 300  
 301      io_thread = None
 302  
 303      def __init__(self, io_thread):
 304          """Initialize the socket."""
 305          self.io_thread = io_thread
 306  
 307      def __getattr__(self, attr):
 308          """Wrap socket attr access for backward-compatibility"""
 309          if attr.startswith("__") and attr.endswith("__"):
 310              # don't wrap magic methods
 311              super().__getattr__(attr)  # type:ignore[misc]
 312          assert self.io_thread is not None
 313          if hasattr(self.io_thread.socket, attr):
 314              warnings.warn(
 315                  f"Accessing zmq Socket attribute {attr} on BackgroundSocket"
 316                  f" is deprecated since ipykernel 4.3.0"
 317                  f" use .io_thread.socket.{attr}",
 318                  DeprecationWarning,
 319                  stacklevel=2,
 320              )
 321              return getattr(self.io_thread.socket, attr)
 322          return super().__getattr__(attr)  # type:ignore[misc]
 323  
 324      def __setattr__(self, attr, value):
 325          """Set an attribute on the socket."""
 326          if attr == "io_thread" or (attr.startswith("__") and attr.endswith("__")):
 327              super().__setattr__(attr, value)
 328          else:
 329              warnings.warn(
 330                  f"Setting zmq Socket attribute {attr} on BackgroundSocket"
 331                  f" is deprecated since ipykernel 4.3.0"
 332                  f" use .io_thread.socket.{attr}",
 333                  DeprecationWarning,
 334                  stacklevel=2,
 335              )
 336              assert self.io_thread is not None
 337              setattr(self.io_thread.socket, attr, value)
 338  
 339      def send(self, msg, *args, **kwargs):
 340          """Send a message to the socket."""
 341          return self.send_multipart([msg], *args, **kwargs)
 342  
 343      def send_multipart(self, *args, **kwargs):
 344          """Schedule send in IO thread"""
 345          assert self.io_thread is not None
 346          return self.io_thread.send_multipart(*args, **kwargs)
 347  
 348  
 349  class OutStream(TextIOBase):
 350      """A file like object that publishes the stream to a 0MQ PUB socket.
 351  
 352      Output is handed off to an IO Thread
 353      """
 354  
 355      # timeout for flush to avoid infinite hang
 356      # in case of misbehavior
 357      flush_timeout = 10
 358      # The time interval between automatic flushes, in seconds.
 359      flush_interval = 0.2
 360      topic = None
 361      encoding = "UTF-8"
 362      _exc: Optional[Any] = None
 363  
 364      def fileno(self):
 365          """
 366          Things like subprocess will peak and write to the fileno() of stderr/stdout.
 367          """
 368          if getattr(self, "_original_stdstream_copy", None) is not None:
 369              return self._original_stdstream_copy
 370          msg = "fileno"
 371          raise io.UnsupportedOperation(msg)
 372  
 373      def _watch_pipe_fd(self):
 374          """
 375          We've redirected standards streams 0 and 1 into a pipe.
 376  
 377          We need to watch in a thread and redirect them to the right places.
 378  
 379          1) the ZMQ channels to show in notebook interfaces,
 380          2) the original stdout/err, to capture errors in terminals.
 381  
 382          We cannot schedule this on the ioloop thread, as this might be blocking.
 383  
 384          """
 385  
 386          try:
 387              bts = os.read(self._fid, PIPE_BUFFER_SIZE)
 388              while bts and self._should_watch:
 389                  self.write(bts.decode(errors="replace"))
 390                  os.write(self._original_stdstream_copy, bts)
 391                  bts = os.read(self._fid, PIPE_BUFFER_SIZE)
 392          except Exception:
 393              self._exc = sys.exc_info()
 394  
 395      def __init__(
 396          self,
 397          session,
 398          pub_thread,
 399          name,
 400          pipe=None,
 401          echo=None,
 402          *,
 403          watchfd=True,
 404          isatty=False,
 405      ):
 406          """
 407          Parameters
 408          ----------
 409          session : object
 410              the session object
 411          pub_thread : threading.Thread
 412              the publication thread
 413          name : str {'stderr', 'stdout'}
 414              the name of the standard stream to replace
 415          pipe : object
 416              the pipe object
 417          echo : bool
 418              whether to echo output
 419          watchfd : bool (default, True)
 420              Watch the file descriptor corresponding to the replaced stream.
 421              This is useful if you know some underlying code will write directly
 422              the file descriptor by its number. It will spawn a watching thread,
 423              that will swap the give file descriptor for a pipe, read from the
 424              pipe, and insert this into the current Stream.
 425          isatty : bool (default, False)
 426              Indication of whether this stream has terminal capabilities (e.g. can handle colors)
 427  
 428          """
 429          if pipe is not None:
 430              warnings.warn(
 431                  "pipe argument to OutStream is deprecated and ignored since ipykernel 4.2.3.",
 432                  DeprecationWarning,
 433                  stacklevel=2,
 434              )
 435          # This is necessary for compatibility with Python built-in streams
 436          self.session = session
 437          if not isinstance(pub_thread, IOPubThread):
 438              # Backward-compat: given socket, not thread. Wrap in a thread.
 439              warnings.warn(
 440                  "Since IPykernel 4.3, OutStream should be created with "
 441                  "IOPubThread, not %r" % pub_thread,
 442                  DeprecationWarning,
 443                  stacklevel=2,
 444              )
 445              pub_thread = IOPubThread(pub_thread)
 446              pub_thread.start()
 447          self.pub_thread = pub_thread
 448          self.name = name
 449          self.topic = b"stream." + name.encode()
 450          self._parent_header: contextvars.ContextVar[Dict[str, Any]] = contextvars.ContextVar(
 451              "parent_header"
 452          )
 453          self._parent_header.set({})
 454          self._thread_to_parent = {}
 455          self._thread_to_parent_header = {}
 456          self._parent_header_global = {}
 457          self._master_pid = os.getpid()
 458          self._flush_pending = False
 459          self._subprocess_flush_pending = False
 460          self._io_loop = pub_thread.io_loop
 461          self._buffer_lock = threading.RLock()
 462          self._buffers = defaultdict(StringIO)
 463          self.echo = None
 464          self._isatty = bool(isatty)
 465          self._should_watch = False
 466          self._local = local()
 467  
 468          if (
 469              watchfd
 470              and (
 471                  (sys.platform.startswith("linux") or sys.platform.startswith("darwin"))
 472                  # Pytest set its own capture. Don't redirect from within pytest.
 473                  and ("PYTEST_CURRENT_TEST" not in os.environ)
 474              )
 475              # allow forcing watchfd (mainly for tests)
 476              or watchfd == "force"
 477          ):
 478              self._should_watch = True
 479              self._setup_stream_redirects(name)
 480  
 481          if echo:
 482              if hasattr(echo, "read") and hasattr(echo, "write"):
 483                  # make sure we aren't trying to echo on the FD we're watching!
 484                  # that would cause an infinite loop, always echoing on itself
 485                  if self._should_watch:
 486                      try:
 487                          echo_fd = echo.fileno()
 488                      except Exception:
 489                          echo_fd = None
 490  
 491                      if echo_fd is not None and echo_fd == self._original_stdstream_fd:
 492                          # echo on the _copy_ we made during
 493                          # this is the actual terminal FD now
 494                          echo = io.TextIOWrapper(
 495                              io.FileIO(
 496                                  self._original_stdstream_copy,
 497                                  "w",
 498                              )
 499                          )
 500                  self.echo = echo
 501              else:
 502                  msg = "echo argument must be a file-like object"
 503                  raise ValueError(msg)
 504  
 505      @property
 506      def parent_header(self):
 507          try:
 508              # asyncio-specific
 509              return self._parent_header.get()
 510          except LookupError:
 511              try:
 512                  # thread-specific
 513                  identity = threading.current_thread().ident
 514                  # retrieve the outermost (oldest ancestor,
 515                  # discounting the kernel thread) thread identity
 516                  while identity in self._thread_to_parent:
 517                      identity = self._thread_to_parent[identity]
 518                  # use the header of the oldest ancestor
 519                  return self._thread_to_parent_header[identity]
 520              except KeyError:
 521                  # global (fallback)
 522                  return self._parent_header_global
 523  
 524      @parent_header.setter
 525      def parent_header(self, value):
 526          self._parent_header_global = value
 527          return self._parent_header.set(value)
 528  
 529      def isatty(self):
 530          """Return a bool indicating whether this is an 'interactive' stream.
 531  
 532          Returns:
 533              Boolean
 534          """
 535          return self._isatty
 536  
 537      def _setup_stream_redirects(self, name):
 538          pr, pw = os.pipe()
 539          fno = self._original_stdstream_fd = getattr(sys, name).fileno()
 540          self._original_stdstream_copy = os.dup(fno)
 541          os.dup2(pw, fno)
 542  
 543          self._fid = pr
 544  
 545          self._exc = None
 546          self.watch_fd_thread = threading.Thread(target=self._watch_pipe_fd)
 547          self.watch_fd_thread.daemon = True
 548          self.watch_fd_thread.start()
 549  
 550      def _is_master_process(self):
 551          return os.getpid() == self._master_pid
 552  
 553      def set_parent(self, parent):
 554          """Set the parent header."""
 555          self.parent_header = extract_header(parent)
 556  
 557      def close(self):
 558          """Close the stream."""
 559          if self._should_watch:
 560              self._should_watch = False
 561              # thread won't wake unless there's something to read
 562              # writing something after _should_watch will not be echoed
 563              os.write(self._original_stdstream_fd, b"\0")
 564              self.watch_fd_thread.join()
 565              # restore original FDs
 566              os.dup2(self._original_stdstream_copy, self._original_stdstream_fd)
 567              os.close(self._original_stdstream_copy)
 568          if self._exc:
 569              etype, value, tb = self._exc
 570              traceback.print_exception(etype, value, tb)
 571          self.pub_thread = None
 572  
 573      @property
 574      def closed(self):
 575          return self.pub_thread is None
 576  
 577      def _schedule_flush(self):
 578          """schedule a flush in the IO thread
 579  
 580          call this on write, to indicate that flush should be called soon.
 581          """
 582          if self._flush_pending:
 583              return
 584          self._flush_pending = True
 585  
 586          # add_timeout has to be handed to the io thread via event pipe
 587          def _schedule_in_thread():
 588              self._io_loop.call_later(self.flush_interval, self._flush)
 589  
 590          self.pub_thread.schedule(_schedule_in_thread)
 591  
 592      def flush(self):
 593          """trigger actual zmq send
 594  
 595          send will happen in the background thread
 596          """
 597          if (
 598              self.pub_thread
 599              and self.pub_thread.thread is not None
 600              and self.pub_thread.thread.is_alive()
 601              and self.pub_thread.thread.ident != threading.current_thread().ident
 602          ):
 603              # request flush on the background thread
 604              self.pub_thread.schedule(self._flush)
 605              # wait for flush to actually get through, if we can.
 606              evt = threading.Event()
 607              self.pub_thread.schedule(evt.set)
 608              # and give a timeout to avoid
 609              if not evt.wait(self.flush_timeout):
 610                  # write directly to __stderr__ instead of warning because
 611                  # if this is happening sys.stderr may be the problem.
 612                  print("IOStream.flush timed out", file=sys.__stderr__)
 613          else:
 614              self._flush()
 615  
 616      def _flush(self):
 617          """This is where the actual send happens.
 618  
 619          _flush should generally be called in the IO thread,
 620          unless the thread has been destroyed (e.g. forked subprocess).
 621          """
 622          self._flush_pending = False
 623          self._subprocess_flush_pending = False
 624  
 625          if self.echo is not None:
 626              try:
 627                  self.echo.flush()
 628              except OSError as e:
 629                  if self.echo is not sys.__stderr__:
 630                      print(f"Flush failed: {e}", file=sys.__stderr__)
 631  
 632          for parent, data in self._flush_buffers():
 633              if data:
 634                  # FIXME: this disables Session's fork-safe check,
 635                  # since pub_thread is itself fork-safe.
 636                  # There should be a better way to do this.
 637                  self.session.pid = os.getpid()
 638                  content = {"name": self.name, "text": data}
 639                  msg = self.session.msg("stream", content, parent=parent)
 640  
 641                  # Each transform either returns a new
 642                  # message or None. If None is returned,
 643                  # the message has been 'used' and we return.
 644                  for hook in self._hooks:
 645                      msg = hook(msg)
 646                      if msg is None:
 647                          return
 648  
 649                  self.session.send(
 650                      self.pub_thread,
 651                      msg,
 652                      ident=self.topic,
 653                  )
 654  
 655      def write(self, string: str) -> Optional[int]:  # type:ignore[override]
 656          """Write to current stream after encoding if necessary
 657  
 658          Returns
 659          -------
 660          len : int
 661              number of items from input parameter written to stream.
 662  
 663          """
 664          parent = self.parent_header
 665  
 666          if not isinstance(string, str):
 667              msg = f"write() argument must be str, not {type(string)}"  # type:ignore[unreachable]
 668              raise TypeError(msg)
 669  
 670          if self.echo is not None:
 671              try:
 672                  self.echo.write(string)
 673              except OSError as e:
 674                  if self.echo is not sys.__stderr__:
 675                      print(f"Write failed: {e}", file=sys.__stderr__)
 676  
 677          if self.pub_thread is None:
 678              msg = "I/O operation on closed file"
 679              raise ValueError(msg)
 680  
 681          is_child = not self._is_master_process()
 682          # only touch the buffer in the IO thread to avoid races
 683          with self._buffer_lock:
 684              self._buffers[frozenset(parent.items())].write(string)
 685          if is_child:
 686              # mp.Pool cannot be trusted to flush promptly (or ever),
 687              # and this helps.
 688              if self._subprocess_flush_pending:
 689                  return None
 690              self._subprocess_flush_pending = True
 691              # We can not rely on self._io_loop.call_later from a subprocess
 692              self.pub_thread.schedule(self._flush)
 693          else:
 694              self._schedule_flush()
 695  
 696          return len(string)
 697  
 698      def writelines(self, sequence):
 699          """Write lines to the stream."""
 700          if self.pub_thread is None:
 701              msg = "I/O operation on closed file"
 702              raise ValueError(msg)
 703          for string in sequence:
 704              self.write(string)
 705  
 706      def writable(self):
 707          """Test whether the stream is writable."""
 708          return True
 709  
 710      def _flush_buffers(self):
 711          """clear the current buffer and return the current buffer data."""
 712          buffers = self._rotate_buffers()
 713          for frozen_parent, buffer in buffers.items():
 714              data = buffer.getvalue()
 715              buffer.close()
 716              yield dict(frozen_parent), data
 717  
 718      def _rotate_buffers(self):
 719          """Returns the current buffer and replaces it with an empty buffer."""
 720          with self._buffer_lock:
 721              old_buffers = self._buffers
 722              self._buffers = defaultdict(StringIO)
 723          return old_buffers
 724  
 725      @property
 726      def _hooks(self):
 727          if not hasattr(self._local, "hooks"):
 728              # create new list for a new thread
 729              self._local.hooks = []
 730          return self._local.hooks
 731  
 732      def register_hook(self, hook):
 733          """
 734          Registers a hook with the thread-local storage.
 735  
 736          Parameters
 737          ----------
 738          hook : Any callable object
 739  
 740          Returns
 741          -------
 742          Either a publishable message, or `None`.
 743          The hook callable must return a message from
 744          the __call__ method if they still require the
 745          `session.send` method to be called after transformation.
 746          Returning `None` will halt that execution path, and
 747          session.send will not be called.
 748          """
 749          self._hooks.append(hook)
 750  
 751      def unregister_hook(self, hook):
 752          """
 753          Un-registers a hook with the thread-local storage.
 754  
 755          Parameters
 756          ----------
 757          hook : Any callable object which has previously been
 758              registered as a hook.
 759  
 760          Returns
 761          -------
 762          bool - `True` if the hook was removed, `False` if it wasn't
 763              found.
 764          """
 765          try:
 766              self._hooks.remove(hook)
 767              return True
 768          except ValueError:
 769              return False
 770