From 5d4888116a69b8f8edb69e51a22b8d4df3a54741 Mon Sep 17 00:00:00 2001
From: Julian Gilbey <julian-git@d-and-j.net>
Date: Tue, 17 Sep 2024 09:48:56 +0100
Subject: [PATCH] Support psutil for finding network addresses (#1033)

* Support psutil

* remove unneeded type: ignore

---------

Co-authored-by: Julian Gilbey <jdg@debian.org>
Co-authored-by: Min RK <benjaminrk@gmail.com>
---
 jupyter_client/localinterfaces.py | 40 +++++++++++++++++++++++++++++--
 1 file changed, 38 insertions(+), 2 deletions(-)

diff --git a/jupyter_client/localinterfaces.py b/jupyter_client/localinterfaces.py
index ca684a6b..8e622f13 100644
--- a/jupyter_client/localinterfaces.py
+++ b/jupyter_client/localinterfaces.py
@@ -151,6 +151,35 @@ def _load_ips_ipconfig() -> None:
     _populate_from_list(addrs)
 
 
+def _load_ips_psutil() -> None:
+    """load ip addresses with netifaces"""
+    import psutil
+
+    global LOCALHOST
+    local_ips = []
+    public_ips = []
+
+    # dict of iface_name: address_list, eg
+    # {"lo": [snicaddr(family=<AddressFamily.AF_INET>, address="127.0.0.1",
+    #   ...), snicaddr(family=<AddressFamily.AF_INET6>, ...)]}
+    for iface, ifaddresses in psutil.net_if_addrs().items():
+        for address_data in ifaddresses:
+            if address_data.family == socket.AF_INET:
+                addr = address_data.address
+                if not (iface.startswith("lo") or addr.startswith("127.")):
+                    public_ips.append(addr)
+                elif not LOCALHOST:
+                    LOCALHOST = addr
+                local_ips.append(addr)
+    if not LOCALHOST:
+        # we never found a loopback interface (can this ever happen?), assume common default
+        LOCALHOST = "127.0.0.1"
+        local_ips.insert(0, LOCALHOST)
+    local_ips.extend(["0.0.0.0", ""])  # noqa
+    LOCAL_IPS[:] = _uniq_stable(local_ips)
+    PUBLIC_IPS[:] = _uniq_stable(public_ips)
+
+
 def _load_ips_netifaces() -> None:
     """load ip addresses with netifaces"""
     import netifaces  # type: ignore[import-not-found]
@@ -227,13 +256,20 @@ def _load_ips(suppress_exceptions: bool = True) -> None:
 
     This function will only ever be called once.
 
-    It will use netifaces to do it quickly if available.
+    If will use psutil to do it quickly if available.
+    If not, it will use netifaces to do it quickly if available.
     Then it will fallback on parsing the output of ifconfig / ip addr / ipconfig, as appropriate.
     Finally, it will fallback on socket.gethostbyname_ex, which can be slow.
     """
 
     try:
-        # first priority, use netifaces
+        # first priority, use psutil
+        try:
+            return _load_ips_psutil()
+        except ImportError:
+            pass
+
+        # second priority, use netifaces
         try:
             return _load_ips_netifaces()
         except ImportError: