refactor ntuple conversion

6 files changed, 30 insertions, 50 deletions

diff --git a/psutil/_common.py b/psutil/_common.py
index 7cc5a205..15e3acf6 100644
--- a/psutil/_common.py
+++ b/psutil/_common.py
@@ -464,6 +464,22 @@ def socktype_to_enum(num):
             return num
 
 
+def conn_to_ntuple(fd, fam, type_, laddr, raddr, status, pid=None):
+    """Convert a raw connection tuple to a proper ntuple."""
+    if fam in (socket.AF_INET, AF_INET6):
+        if laddr:
+            laddr = addr(*laddr)
+        if raddr:
+            raddr = addr(*raddr)
+    fam = sockfam_to_enum(fam)
+    type_ = socktype_to_enum(type_)
+    if pid is None:
+        nt = pconn(fd, fam, type_, laddr, raddr, status)
+    else:
+        nt = sconn(fd, fam, type_, laddr, raddr, status, pid)
+    return nt
+
+
 def deprecated_method(replacement):
     """A decorator which can be used to mark a method as deprecated
     'replcement' is the method name which will be called instead.
diff --git a/psutil/_psaix.py b/psutil/_psaix.py
index b24325d1..0542d13c 100644
--- a/psutil/_psaix.py
+++ b/psutil/_psaix.py
@@ -220,6 +220,7 @@ def net_connections(kind, _pid=-1):
                          % (kind, ', '.join([repr(x) for x in cmap])))
     families, types = _common.conn_tmap[kind]
     rawlist = cext.net_connections(_pid)
+    # XXX: why is this not a list?
     ret = set()
     for item in rawlist:
         fd, fam, type_, laddr, raddr, status, pid = item
diff --git a/psutil/_psbsd.py b/psutil/_psbsd.py
index e90bdc84..e32d1289 100644
--- a/psutil/_psbsd.py
+++ b/psutil/_psbsd.py
@@ -10,22 +10,18 @@ import functools
 import os
 import xml.etree.ElementTree as ET
 from collections import namedtuple
-from socket import AF_INET
 from collections import defaultdict
 
 from . import _common
 from . import _psposix
 from . import _psutil_bsd as cext
 from . import _psutil_posix as cext_posix
-from ._common import AF_INET6
 from ._common import conn_tmap
 from ._common import FREEBSD
 from ._common import memoize
 from ._common import memoize_when_activated
 from ._common import NETBSD
 from ._common import OPENBSD
-from ._common import sockfam_to_enum
-from ._common import socktype_to_enum
 from ._common import usage_percent
 from ._compat import which
 
@@ -407,14 +403,8 @@ def net_connections(kind):
                 # have a very short lifetime so maybe the kernel
                 # can't initialize their status?
                 status = TCP_STATUSES[cext.PSUTIL_CONN_NONE]
-            if fam in (AF_INET, AF_INET6):
-                if laddr:
-                    laddr = _common.addr(*laddr)
-                if raddr:
-                    raddr = _common.addr(*raddr)
-            fam = sockfam_to_enum(fam)
-            type = socktype_to_enum(type)
-            nt = _common.sconn(fd, fam, type, laddr, raddr, status, pid)
+            nt = _common.conn_to_ntuple(fd, fam, type, laddr, raddr, status,
+                                        pid)
             ret.add(nt)
     return list(ret)
 
@@ -769,6 +759,7 @@ class Process(object):
 
         if NETBSD:
             families, types = conn_tmap[kind]
+            # XXX: why is this not a list?
             ret = set()
             rawlist = cext.net_connections(self.pid)
             for item in rawlist:
@@ -779,14 +770,8 @@ class Process(object):
                         status = TCP_STATUSES[status]
                     except KeyError:
                         status = TCP_STATUSES[cext.PSUTIL_CONN_NONE]
-                    if fam in (AF_INET, AF_INET6):
-                        if laddr:
-                            laddr = _common.addr(*laddr)
-                        if raddr:
-                            raddr = _common.addr(*raddr)
-                    fam = sockfam_to_enum(fam)
-                    type = socktype_to_enum(type)
-                    nt = _common.pconn(fd, fam, type, laddr, raddr, status)
+                    nt = _common.conn_to_ntuple(
+                        fd, fam, type, laddr, raddr, status)
                     ret.add(nt)
             self._assert_alive()
             return list(ret)
@@ -796,18 +781,13 @@ class Process(object):
         ret = []
         for item in rawlist:
             fd, fam, type, laddr, raddr, status = item
-            if fam in (AF_INET, AF_INET6):
-                if laddr:
-                    laddr = _common.addr(*laddr)
-                if raddr:
-                    raddr = _common.addr(*raddr)
-            fam = sockfam_to_enum(fam)
-            type = socktype_to_enum(type)
             status = TCP_STATUSES[status]
-            nt = _common.pconn(fd, fam, type, laddr, raddr, status)
+            nt = _common.conn_to_ntuple(fd, fam, type, laddr, raddr, status)
             ret.append(nt)
+
         if OPENBSD:
             self._assert_alive()
+
         return ret
 
     @wrap_exceptions
diff --git a/psutil/_psosx.py b/psutil/_psosx.py
index 7459a0f3..7ef9326d 100644
--- a/psutil/_psosx.py
+++ b/psutil/_psosx.py
@@ -8,20 +8,16 @@ import contextlib
 import errno
 import functools
 import os
-from socket import AF_INET
 from collections import namedtuple
 
 from . import _common
 from . import _psposix
 from . import _psutil_osx as cext
 from . import _psutil_posix as cext_posix
-from ._common import AF_INET6
 from ._common import conn_tmap
 from ._common import isfile_strict
 from ._common import memoize_when_activated
 from ._common import parse_environ_block
-from ._common import sockfam_to_enum
-from ._common import socktype_to_enum
 from ._common import usage_percent
 
 
@@ -530,14 +526,7 @@ class Process(object):
         for item in rawlist:
             fd, fam, type, laddr, raddr, status = item
             status = TCP_STATUSES[status]
-            fam = sockfam_to_enum(fam)
-            type = socktype_to_enum(type)
-            if fam in (AF_INET, AF_INET6):
-                if laddr:
-                    laddr = _common.addr(*laddr)
-                if raddr:
-                    raddr = _common.addr(*raddr)
-            nt = _common.pconn(fd, fam, type, laddr, raddr, status)
+            nt = _common.conn_to_ntuple(fd, fam, type, laddr, raddr, status)
             ret.append(nt)
         return ret
 
diff --git a/psutil/_pssunos.py b/psutil/_pssunos.py
index 6d7fda85..07298a73 100644
--- a/psutil/_pssunos.py
+++ b/psutil/_pssunos.py
@@ -262,6 +262,7 @@ def net_connections(kind, _pid=-1):
             continue
         if type_ not in types:
             continue
+        # TODO: refactor and use _common.conn_to_ntuple.
         if fam in (AF_INET, AF_INET6):
             if laddr:
                 laddr = _common.addr(*laddr)
diff --git a/psutil/_pswindows.py b/psutil/_pswindows.py
index 3f131980..891891a8 100644
--- a/psutil/_pswindows.py
+++ b/psutil/_pswindows.py
@@ -39,8 +39,6 @@ from ._common import isfile_strict
 from ._common import memoize
 from ._common import memoize_when_activated
 from ._common import parse_environ_block
-from ._common import sockfam_to_enum
-from ._common import socktype_to_enum
 from ._common import usage_percent
 from ._compat import long
 from ._compat import lru_cache
@@ -388,17 +386,12 @@ def net_connections(kind, _pid=-1):
     ret = set()
     for item in rawlist:
         fd, fam, type, laddr, raddr, status, pid = item
-        if laddr:
-            laddr = _common.addr(*laddr)
-        if raddr:
-            raddr = _common.addr(*raddr)
         status = TCP_STATUSES[status]
-        fam = sockfam_to_enum(fam)
-        type = socktype_to_enum(type)
         if _pid == -1:
-            nt = _common.sconn(fd, fam, type, laddr, raddr, status, pid)
+            nt = _common.conn_to_ntuple(fd, fam, type, laddr, raddr, status,
+                                        pid=pid)
         else:
-            nt = _common.pconn(fd, fam, type, laddr, raddr, status)
+            nt = _common.conn_to_ntuple(fd, fam, type, laddr, raddr, status)
         ret.add(nt)
     return list(ret)