pyi_splash.py 7.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211
  1. # -----------------------------------------------------------------------------
  2. # Copyright (c) 2005-2023, PyInstaller Development Team.
  3. #
  4. # Licensed under the Apache License, Version 2.0 (the "License");
  5. # you may not use this file except in compliance with the License.
  6. #
  7. # The full license is in the file COPYING.txt, distributed with this software.
  8. #
  9. # SPDX-License-Identifier: Apache-2.0
  10. # -----------------------------------------------------------------------------
  11. # This module is not a "fake module" in the classical sense, but a real module that can be imported. It acts as an RPC
  12. # interface for the functions of the bootloader.
  13. """
  14. This module connects to the bootloader to send messages to the splash screen.
  15. It is intended to act as a RPC interface for the functions provided by the bootloader, such as displaying text or
  16. closing. This makes the users python program independent of how the communication with the bootloader is implemented,
  17. since a consistent API is provided.
  18. To connect to the bootloader, it connects to a local tcp socket whose port is passed through the environment variable
  19. '_PYI_SPLASH_IPC'. The bootloader creates a server socket and accepts every connection request. Since the os-module,
  20. which is needed to request the environment variable, is not available at boot time, the module does not establish the
  21. connection until initialization.
  22. The protocol by which the Python interpreter communicates with the bootloader is implemented in this module.
  23. This module does not support reloads while the splash screen is displayed, i.e. it cannot be reloaded (such as by
  24. importlib.reload), because the splash screen closes automatically when the connection to this instance of the module
  25. is lost.
  26. """
  27. import atexit
  28. import os
  29. # Import the _socket module instead of the socket module. All used functions to connect to the ipc system are
  30. # provided by the C module and the users program does not necessarily need to include the socket module and all
  31. # required modules it uses.
  32. import _socket
  33. __all__ = ["CLOSE_CONNECTION", "FLUSH_CHARACTER", "is_alive", "close", "update_text"]
  34. try:
  35. # The user might have excluded logging from imports.
  36. import logging as _logging
  37. except ImportError:
  38. _logging = None
  39. try:
  40. # The user might have excluded functools from imports.
  41. from functools import update_wrapper
  42. except ImportError:
  43. update_wrapper = None
  44. # Utility
  45. def _log(level, msg, *args, **kwargs):
  46. """
  47. Conditional wrapper around logging module. If the user excluded logging from the imports or it was not imported,
  48. this function should handle it and avoid using the logger.
  49. """
  50. if _logging:
  51. logger = _logging.getLogger(__name__)
  52. logger.log(level, msg, *args, **kwargs)
  53. # These constants define single characters which are needed to send commands to the bootloader. Those constants are
  54. # also set in the tcl script.
  55. CLOSE_CONNECTION = b'\x04' # ASCII End-of-Transmission character
  56. FLUSH_CHARACTER = b'\x0D' # ASCII Carriage Return character
  57. # Module internal variables
  58. _initialized = False
  59. # Keep these variables always synchronized
  60. _ipc_socket_closed = True
  61. _ipc_socket = _socket.socket(_socket.AF_INET, _socket.SOCK_STREAM)
  62. def _initialize():
  63. """
  64. Initialize this module
  65. :return:
  66. """
  67. global _initialized, _ipc_socket_closed
  68. # If _ipc_port is zero, the splash screen is intentionally suppressed (for example, we are in sub-process spawned
  69. # via sys.executable). Mark the splash screen as initialized, but do not attempt to connect.
  70. if _ipc_port == 0:
  71. _initialized = True
  72. return
  73. # Attempt to connect to the splash screen process.
  74. try:
  75. _ipc_socket.connect(("127.0.0.1", _ipc_port))
  76. _ipc_socket_closed = False
  77. _initialized = True
  78. _log(10, "IPC connection to the splash screen was successfully established.") # log-level: debug
  79. except OSError as err:
  80. raise ConnectionError(f"Could not connect to TCP port {_ipc_port}.") from err
  81. # We expect a splash screen from the bootloader, but if _PYI_SPLASH_IPC is not set, the module cannot connect to it.
  82. # _PYI_SPLASH_IPC being set to zero indicates that splash screen should be (gracefully) suppressed; i.e., the calls
  83. # in this module should become no-op without generating warning messages.
  84. try:
  85. _ipc_port = int(os.environ['_PYI_SPLASH_IPC'])
  86. del os.environ['_PYI_SPLASH_IPC']
  87. # Initialize the connection upon importing this module. This will establish a connection to the bootloader's TCP
  88. # server socket.
  89. _initialize()
  90. except (KeyError, ValueError):
  91. # log-level: warning
  92. _log(
  93. 30,
  94. "The environment does not allow connecting to the splash screen. Did bootloader fail to initialize it?",
  95. exc_info=True,
  96. )
  97. except ConnectionError:
  98. # log-level: error
  99. _log(40, "Failed to connect to the bootloader's IPC server!", exc_info=True)
  100. def _check_connection(func):
  101. """
  102. Utility decorator for checking whether the function should be executed.
  103. The wrapped function may raise a ConnectionError if the module was not initialized correctly.
  104. """
  105. def wrapper(*args, **kwargs):
  106. """
  107. Executes the wrapped function if the environment allows it.
  108. That is, if the connection to to bootloader has not been closed and the module is initialized.
  109. :raises RuntimeError: if the module was not initialized correctly.
  110. """
  111. if _initialized and _ipc_socket_closed:
  112. if _ipc_port != 0:
  113. _log(10, "Connection to splash screen has already been closed.") # log-level: debug
  114. return
  115. elif not _initialized:
  116. raise RuntimeError("This module is not initialized; did it fail to load?")
  117. return func(*args, **kwargs)
  118. if update_wrapper:
  119. # For runtime introspection
  120. update_wrapper(wrapper, func)
  121. return wrapper
  122. @_check_connection
  123. def _send_command(cmd, args=None):
  124. """
  125. Send the command followed by args to the splash screen.
  126. :param str cmd: The command to send. All command have to be defined as procedures in the tcl splash screen script.
  127. :param list[str] args: All arguments to send to the receiving function
  128. """
  129. if args is None:
  130. args = []
  131. full_cmd = "%s(%s)" % (cmd, " ".join(args))
  132. try:
  133. _ipc_socket.sendall(full_cmd.encode("utf-8") + FLUSH_CHARACTER)
  134. except OSError as err:
  135. raise ConnectionError(f"Unable to send command {full_cmd!r} to the bootloader") from err
  136. def is_alive():
  137. """
  138. Indicates whether the module can be used.
  139. Returns False if the module is either not initialized or was disabled by closing the splash screen. Otherwise,
  140. the module should be usable.
  141. """
  142. return _initialized and not _ipc_socket_closed
  143. @_check_connection
  144. def update_text(msg: str):
  145. """
  146. Updates the text on the splash screen window.
  147. :param str msg: the text to be displayed
  148. :raises ConnectionError: If the OS fails to write to the socket.
  149. :raises RuntimeError: If the module is not initialized.
  150. """
  151. _send_command("update_text", [msg])
  152. def close():
  153. """
  154. Close the connection to the ipc tcp server socket.
  155. This will close the splash screen and renders this module unusable. After this function is called, no connection
  156. can be opened to the splash screen again and all functions in this module become unusable.
  157. """
  158. global _ipc_socket_closed
  159. if _initialized and not _ipc_socket_closed:
  160. _ipc_socket.sendall(CLOSE_CONNECTION)
  161. _ipc_socket.close()
  162. _ipc_socket_closed = True
  163. @atexit.register
  164. def _exit():
  165. close()