From 5d4be42bd309dc3c7e93d5d319577ec4b8027445 Mon Sep 17 00:00:00 2001 From: Giampaolo Rodola Date: Fri, 20 Dec 2024 15:20:48 +0100 Subject: [PATCH] refact: check valid net_connections() kinds from a single point --- psutil/__init__.py | 11 ++++++++++- psutil/_common.py | 2 +- psutil/_psaix.py | 6 ------ psutil/_psbsd.py | 11 ----------- psutil/_pslinux.py | 5 ----- psutil/_psosx.py | 5 ----- psutil/_pssunos.py | 8 -------- psutil/_pswindows.py | 5 ----- 8 files changed, 11 insertions(+), 42 deletions(-) diff --git a/psutil/__init__.py b/psutil/__init__.py index 490e7b394..c7d377f60 100644 --- a/psutil/__init__.py +++ b/psutil/__init__.py @@ -18,7 +18,6 @@ Supported Python versions are cPython 3.6+ and PyPy. """ - import collections import contextlib import datetime @@ -275,6 +274,14 @@ def _pprint_secs(secs): return datetime.datetime.fromtimestamp(secs).strftime(fmt) +def _check_conn_kind(kind): + """Check net_connections()'s `kind` parameter.""" + kinds = tuple(_common.conn_tmap) + if kind not in kinds: + msg = f"invalid kind argument {kind!r}; valid ones are: {kinds}" + raise ValueError(msg) + + # ===================================================================== # --- Process class # ===================================================================== @@ -1231,6 +1238,7 @@ def net_connections(self, kind='inet'): | all | the sum of all the possible families and protocols | +------------+----------------------------------------------------+ """ + _check_conn_kind(kind) return self._proc.net_connections(kind) @_common.deprecated_method(replacement="net_connections") @@ -2191,6 +2199,7 @@ def net_connections(kind='inet'): On macOS this function requires root privileges. """ + _check_conn_kind(kind) return _psplatform.net_connections(kind) diff --git a/psutil/_common.py b/psutil/_common.py index ffa6dcef3..a0986ca7b 100644 --- a/psutil/_common.py +++ b/psutil/_common.py @@ -249,7 +249,7 @@ class BatteryTime(enum.IntEnum): "udp6": ([AF_INET6], [SOCK_DGRAM]), }) -if AF_UNIX is not None: +if AF_UNIX is not None and not SUNOS: conn_tmap.update({"unix": ([AF_UNIX], [SOCK_STREAM, SOCK_DGRAM])}) diff --git a/psutil/_psaix.py b/psutil/_psaix.py index 94a07a3a5..b774d8d42 100644 --- a/psutil/_psaix.py +++ b/psutil/_psaix.py @@ -204,12 +204,6 @@ def net_connections(kind, _pid=-1): """Return socket connections. If pid == -1 return system-wide connections (as opposed to connections opened by one process only). """ - cmap = _common.conn_tmap - if kind not in cmap: - raise ValueError( - "invalid %r kind argument; choose between %s" - % (kind, ', '.join([repr(x) for x in cmap])) - ) families, types = _common.conn_tmap[kind] rawlist = cext.net_connections(_pid) ret = [] diff --git a/psutil/_psbsd.py b/psutil/_psbsd.py index 209e5476c..0628facbe 100644 --- a/psutil/_psbsd.py +++ b/psutil/_psbsd.py @@ -435,14 +435,8 @@ def net_if_stats(): def net_connections(kind): """System-wide network connections.""" - if kind not in _common.conn_tmap: - raise ValueError( - "invalid %r kind argument; choose between %s" - % (kind, ', '.join([repr(x) for x in conn_tmap])) - ) families, types = conn_tmap[kind] ret = set() - if OPENBSD: rawlist = cext.net_connections(-1, families, types) elif NETBSD: @@ -820,11 +814,6 @@ def threads(self): @wrap_exceptions def net_connections(self, kind='inet'): - if kind not in conn_tmap: - raise ValueError( - "invalid %r kind argument; choose between %s" - % (kind, ', '.join([repr(x) for x in conn_tmap])) - ) families, types = conn_tmap[kind] ret = [] diff --git a/psutil/_pslinux.py b/psutil/_pslinux.py index 5d5595feb..b67964cf6 100644 --- a/psutil/_pslinux.py +++ b/psutil/_pslinux.py @@ -955,11 +955,6 @@ def process_unix(file, family, inodes, filter_pid=None): yield (fd, family, type_, path, raddr, status, pid) def retrieve(self, kind, pid=None): - if kind not in self.tmap: - raise ValueError( - "invalid %r kind argument; choose between %s" - % (kind, ', '.join([repr(x) for x in self.tmap])) - ) self._procfs_path = get_procfs_path() if pid is not None: inodes = self.get_proc_inodes(pid) diff --git a/psutil/_psosx.py b/psutil/_psosx.py index c6078ea3d..3fdbcd6d7 100644 --- a/psutil/_psosx.py +++ b/psutil/_psosx.py @@ -500,11 +500,6 @@ def open_files(self): @wrap_exceptions def net_connections(self, kind='inet'): - if kind not in conn_tmap: - raise ValueError( - "invalid %r kind argument; choose between %s" - % (kind, ', '.join([repr(x) for x in conn_tmap])) - ) families, types = conn_tmap[kind] rawlist = cext.proc_net_connections(self.pid, families, types) ret = [] diff --git a/psutil/_pssunos.py b/psutil/_pssunos.py index 09a79fef4..b741792f3 100644 --- a/psutil/_pssunos.py +++ b/psutil/_pssunos.py @@ -260,14 +260,6 @@ def net_connections(kind, _pid=-1): connections (as opposed to connections opened by one process only). Only INET sockets are returned (UNIX are not). """ - cmap = _common.conn_tmap.copy() - if _pid == -1: - cmap.pop('unix', 0) - if kind not in cmap: - raise ValueError( - "invalid %r kind argument; choose between %s" - % (kind, ', '.join([repr(x) for x in cmap])) - ) families, types = _common.conn_tmap[kind] rawlist = cext.net_connections(_pid) ret = set() diff --git a/psutil/_pswindows.py b/psutil/_pswindows.py index be100493e..e3c7db3ca 100644 --- a/psutil/_pswindows.py +++ b/psutil/_pswindows.py @@ -349,11 +349,6 @@ def net_connections(kind, _pid=-1): """Return socket connections. If pid == -1 return system-wide connections (as opposed to connections opened by one process only). """ - if kind not in conn_tmap: - raise ValueError( - "invalid %r kind argument; choose between %s" - % (kind, ', '.join([repr(x) for x in conn_tmap])) - ) families, types = conn_tmap[kind] rawlist = cext.net_connections(_pid, families, types) ret = set()