Selaa lähdekoodia

Merge pull request #625 from tobb10001/fix/412-web-browser-opens-to-early

Wait for the port to be open before opening the web browser.
Rodja Trappe 2 vuotta sitten
vanhempi
säilyke
2ec46e4447
3 muutettua tiedostoa jossa 126 lisäystä ja 4 poistoa
  1. 48 1
      nicegui/helpers.py
  2. 2 3
      nicegui/run.py
  3. 76 0
      tests/test_helpers.py

+ 48 - 1
nicegui/helpers.py

@@ -1,8 +1,12 @@
 import asyncio
 import functools
 import inspect
+import socket
+import threading
+import time
+import webbrowser
 from contextlib import nullcontext
-from typing import TYPE_CHECKING, Any, Awaitable, Callable, Optional, Union
+from typing import TYPE_CHECKING, Any, Awaitable, Callable, Optional, Tuple, Union
 
 from . import background_tasks, globals
 
@@ -33,3 +37,46 @@ def safe_invoke(func: Union[Callable, Awaitable], client: Optional['Client'] = N
                 background_tasks.create(result_with_client())
     except Exception as e:
         globals.handle_exception(e)
+
+
+def is_port_open(host: str, port: int) -> bool:
+    """Check if the port is open by checking if a TCP connection can be established."""
+    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    try:
+        sock.connect((host, port))
+    except (ConnectionRefusedError, TimeoutError):
+        return False
+    else:
+        return True
+    finally:
+        sock.close()
+
+
+def schedule_browser(host: str, port: int) -> Tuple[threading.Thread, threading.Event]:
+    """Wait non-blockingly for the port to be open, then start a webbrowser.
+
+    This function launches a thread in order to be non-blocking. This thread then uses
+    `is_port_open` to check when the port opens. When connectivity is confirmed, the
+    webbrowser is launched using `webbrowser.open`
+
+    The thread is created as a daemon thread, in order to not interfere with Ctrl+C.
+
+    If you need to stop this thread, you can do this by setting the Event, that gets
+    returned. The thread will stop with the next loop without opening the browser.
+
+    :return: A tuple consisting of the actual thread object and an event for stopping
+        the thread.
+    """
+    cancel = threading.Event()
+
+    def in_thread(host: str, port: int) -> None:
+        while not is_port_open(host, port):
+            if cancel.is_set():
+                return
+            time.sleep(0.1)
+        webbrowser.open(f'http://{host}:{port}/')
+
+    host = host if host != "0.0.0.0" else "127.0.0.1"
+    thread = threading.Thread(target=in_thread, args=(host, port), daemon=True)
+    thread.start()
+    return thread, cancel

+ 2 - 3
nicegui/run.py

@@ -2,14 +2,13 @@ import logging
 import multiprocessing
 import os
 import sys
-import webbrowser
 from typing import List, Optional, Tuple
 
 import uvicorn
 from uvicorn.main import STARTUP_FAILURE
 from uvicorn.supervisors import ChangeReload, Multiprocess
 
-from . import globals, native_mode
+from . import globals, helpers, native_mode
 
 
 def run(*,
@@ -88,7 +87,7 @@ def run(*,
     os.environ['NICEGUI_URL'] = f'http://{host}:{port}'
 
     if show:
-        webbrowser.open(f'http://{host if host != "0.0.0.0" else "127.0.0.1"}:{port}/')
+        helpers.schedule_browser(host, port)
 
     def split_args(args: str) -> List[str]:
         return [a.strip() for a in args.split(',')]

+ 76 - 0
tests/test_helpers.py

@@ -0,0 +1,76 @@
+import contextlib
+import socket
+import time
+import webbrowser
+
+from nicegui import helpers
+
+
+def test_is_port_open_when_open():
+    with contextlib.closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as sock:
+        sock.bind(('127.0.0.1', 0))  # port = 0 => the OS chooses a port for us
+        sock.listen(1)
+        host, port = sock.getsockname()
+    assert not helpers.is_port_open(host, port), 'after closing the socket, the port should be free'
+
+
+def test_is_port_open_when_closed():
+    with contextlib.closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as sock:
+        sock.bind(('127.0.0.1', port))
+        sock.listen(1)
+        assert helpers.is_port_open(host, port), 'after opening the socket, the port should be detected'
+
+
+def test_schedule_browser(monkeypatch):
+
+    called_with_url = None
+
+    def mock_webbrowser_open(url):
+        nonlocal called_with_url
+        called_with_url = url
+
+    monkeypatch.setattr(webbrowser, "open", mock_webbrowser_open)
+
+    with contextlib.closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as sock:
+
+        sock.bind(('127.0.0.1', 0))
+        host, port = sock.getsockname()
+
+        thread, cancel_event = helpers.schedule_browser(host, port)
+
+        try:
+            # port bound, but not opened yet
+            assert called_with_url is None
+
+            sock.listen()
+            # port opened
+            time.sleep(1)
+            assert called_with_url == f"http://{host}:{port}/"
+        finally:
+            cancel_event.set()
+
+
+def test_canceling_schedule_browser(monkeypatch):
+
+    called_with_url = None
+
+    def mock_webbrowser_open(url):
+        nonlocal called_with_url
+        called_with_url = url
+
+    monkeypatch.setattr(webbrowser, "open", mock_webbrowser_open)
+
+    # find a free port ...
+    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    sock.bind(('127.0.0.1', 0))
+    sock.listen(1)
+    host, port = sock.getsockname()
+    # ... and close it so schedule_browser does not launch the browser
+    sock.close()
+
+    thread, cancel_event = helpers.schedule_browser(host, port)
+    time.sleep(0.2)
+    cancel_event.set()
+    time.sleep(0.2)
+    assert not thread.is_alive()
+    assert called_with_url is None