fix #582: turn *connections() constants into enums on python 3.4+

9 files changed, 71 insertions, 21 deletions

diff --git a/HISTORY.rst b/HISTORY.rst
index 4632cca5..b6a19232 100644
--- a/HISTORY.rst
+++ b/HISTORY.rst
@@ -6,6 +6,8 @@ Bug tracker at https://github.com/giampaolo/psutil/issues
 **Enhancements**
 
 - #581: add .gitignore. (patch by Gabi Davar)
+- #582: connection constants returned by psutil.net_connections() and
+  psutil.Process.connections() were turned from int to enums on Python > 3.4.
 
 
 2.2.1 - 2015-02-02
diff --git a/README.rst b/README.rst
index e0e1c367..12843b8d 100644
--- a/README.rst
+++ b/README.rst
@@ -142,10 +142,10 @@ Network
      'lo': netio(bytes_sent=2838627, bytes_recv=2838627, packets_sent=30567, packets_recv=30567, errin=0, errout=0, dropin=0, dropout=0)}
     >>>
     >>> psutil.net_connections()
-    [pconn(fd=115, family=2, type=1, laddr=('10.0.0.1', 48776), raddr=('93.186.135.91', 80), status='ESTABLISHED', pid=1254),
-     pconn(fd=117, family=2, type=1, laddr=('10.0.0.1', 43761), raddr=('72.14.234.100', 80), status='CLOSING', pid=2987),
-     pconn(fd=-1, family=2, type=1, laddr=('10.0.0.1', 60759), raddr=('72.14.234.104', 80), status='ESTABLISHED', pid=None),
-     pconn(fd=-1, family=2, type=1, laddr=('10.0.0.1', 51314), raddr=('72.14.234.83', 443), status='SYN_SENT', pid=None)
+    [pconn(fd=115, family=<AddressFamily.AF_INET: 2>, type=<SocketType.SOCK_STREAM: 1>, laddr=('10.0.0.1', 48776), raddr=('93.186.135.91', 80), status='ESTABLISHED', pid=1254),
+     pconn(fd=117, family=<AddressFamily.AF_INET: 2>, type=<SocketType.SOCK_STREAM: 1>, laddr=('10.0.0.1', 43761), raddr=('72.14.234.100', 80), status='CLOSING', pid=2987),
+     pconn(fd=-1, family=<AddressFamily.AF_INET: 2>, type=<SocketType.SOCK_STREAM: 1>, laddr=('10.0.0.1', 60759), raddr=('72.14.234.104', 80), status='ESTABLISHED', pid=None),
+     pconn(fd=-1, family=<AddressFamily.AF_INET: 2>, type=<SocketType.SOCK_STREAM: 1>, laddr=('10.0.0.1', 51314), raddr=('72.14.234.83', 443), status='SYN_SENT', pid=None)
      ...]
 
 Other system info
@@ -229,10 +229,10 @@ Process management
     [popenfile(path='/home/giampaolo/svn/psutil/somefile', fd=3)]
     >>>
     >>> p.connections()
-    [pconn(fd=115, family=2, type=1, laddr=('10.0.0.1', 48776), raddr=('93.186.135.91', 80), status='ESTABLISHED'),
-     pconn(fd=117, family=2, type=1, laddr=('10.0.0.1', 43761), raddr=('72.14.234.100', 80), status='CLOSING'),
-     pconn(fd=119, family=2, type=1, laddr=('10.0.0.1', 60759), raddr=('72.14.234.104', 80), status='ESTABLISHED'),
-     pconn(fd=123, family=2, type=1, laddr=('10.0.0.1', 51314), raddr=('72.14.234.83', 443), status='SYN_SENT')]
+    [pconn(fd=115, family=<AddressFamily.AF_INET: 2>, type=<SocketType.SOCK_STREAM: 1>, laddr=('10.0.0.1', 48776), raddr=('93.186.135.91', 80), status='ESTABLISHED'),
+     pconn(fd=117, family=<AddressFamily.AF_INET: 2>, type=<SocketType.SOCK_STREAM: 1>, laddr=('10.0.0.1', 43761), raddr=('72.14.234.100', 80), status='CLOSING'),
+     pconn(fd=119, family=<AddressFamily.AF_INET: 2>, type=<SocketType.SOCK_STREAM: 1>, laddr=('10.0.0.1', 60759), raddr=('72.14.234.104', 80), status='ESTABLISHED'),
+     pconn(fd=123, family=<AddressFamily.AF_INET: 2>, type=<SocketType.SOCK_STREAM: 1>, laddr=('10.0.0.1', 51314), raddr=('72.14.234.83', 443), status='SYN_SENT')]
     >>>
     >>> p.num_threads()
     4
diff --git a/docs/index.rst b/docs/index.rst
index e6497822..86651b41 100644
--- a/docs/index.rst
+++ b/docs/index.rst
@@ -391,10 +391,10 @@ Network
 
     >>> import psutil
     >>> psutil.net_connections()
-    [pconn(fd=115, family=2, type=1, laddr=('10.0.0.1', 48776), raddr=('93.186.135.91', 80), status='ESTABLISHED', pid=1254),
-     pconn(fd=117, family=2, type=1, laddr=('10.0.0.1', 43761), raddr=('72.14.234.100', 80), status='CLOSING', pid=2987),
-     pconn(fd=-1, family=2, type=1, laddr=('10.0.0.1', 60759), raddr=('72.14.234.104', 80), status='ESTABLISHED', pid=None),
-     pconn(fd=-1, family=2, type=1, laddr=('10.0.0.1', 51314), raddr=('72.14.234.83', 443), status='SYN_SENT', pid=None)
+    [pconn(fd=115, family=<AddressFamily.AF_INET: 2>, type=<SocketType.SOCK_STREAM: 1>, laddr=('10.0.0.1', 48776), raddr=('93.186.135.91', 80), status='ESTABLISHED', pid=1254),
+     pconn(fd=117, family=<AddressFamily.AF_INET: 2>, type=<SocketType.SOCK_STREAM: 1>, laddr=('10.0.0.1', 43761), raddr=('72.14.234.100', 80), status='CLOSING', pid=2987),
+     pconn(fd=-1, family=<AddressFamily.AF_INET: 2>, type=<SocketType.SOCK_STREAM: 1>, laddr=('10.0.0.1', 60759), raddr=('72.14.234.104', 80), status='ESTABLISHED', pid=None),
+     pconn(fd=-1, family=<AddressFamily.AF_INET: 2>, type=<SocketType.SOCK_STREAM: 1>, laddr=('10.0.0.1', 51314), raddr=('72.14.234.83', 443), status='SYN_SENT', pid=None)
      ...]
 
   .. note:: (OSX) :class:`psutil.AccessDenied` is always raised unless running
@@ -1041,10 +1041,10 @@ Process class
       >>> p.name()
       'firefox'
       >>> p.connections()
-      [pconn(fd=115, family=2, type=1, laddr=('10.0.0.1', 48776), raddr=('93.186.135.91', 80), status='ESTABLISHED'),
-       pconn(fd=117, family=2, type=1, laddr=('10.0.0.1', 43761), raddr=('72.14.234.100', 80), status='CLOSING'),
-       pconn(fd=119, family=2, type=1, laddr=('10.0.0.1', 60759), raddr=('72.14.234.104', 80), status='ESTABLISHED'),
-       pconn(fd=123, family=2, type=1, laddr=('10.0.0.1', 51314), raddr=('72.14.234.83', 443), status='SYN_SENT')]
+      [pconn(fd=115, family=<AddressFamily.AF_INET: 2>, type=<SocketType.SOCK_STREAM: 1>, laddr=('10.0.0.1', 48776), raddr=('93.186.135.91', 80), status='ESTABLISHED'),
+       pconn(fd=117, family=<AddressFamily.AF_INET: 2>, type=<SocketType.SOCK_STREAM: 1>, laddr=('10.0.0.1', 43761), raddr=('72.14.234.100', 80), status='CLOSING'),
+       pconn(fd=119, family=<AddressFamily.AF_INET: 2>, type=<SocketType.SOCK_STREAM: 1>, laddr=('10.0.0.1', 60759), raddr=('72.14.234.104', 80), status='ESTABLISHED'),
+       pconn(fd=123, family=<AddressFamily.AF_INET: 2>, type=<SocketType.SOCK_STREAM: 1>, laddr=('10.0.0.1', 51314), raddr=('72.14.234.83', 443), status='SYN_SENT')]
 
   .. method:: is_running()
 
diff --git a/psutil/_common.py b/psutil/_common.py
index 92d0fd0c..07f362a9 100644
--- a/psutil/_common.py
+++ b/psutil/_common.py
@@ -17,6 +17,10 @@ try:
     import threading
 except ImportError:
     import dummy_threading as threading
+try:
+    import enum  # py >= 3.4
+except ImportError:
+    enum = None
 
 from collections import namedtuple
 from socket import AF_INET, SOCK_STREAM, SOCK_DGRAM
@@ -160,6 +164,30 @@ def isfile_strict(path):
         return stat.S_ISREG(st.st_mode)
 
 
+def sockfam_to_enum(num):
+    """Convert a numeric socket family value to an IntEnum member.
+    If it's not a known member, return the numeric value itself.
+    """
+    if enum is None:
+        return num
+    try:
+        return socket.AddressFamily(num)
+    except ValueError:
+        return num
+
+
+def socktype_to_enum(num):
+    """Convert a numeric socket type value to an IntEnum member.
+    If it's not a known member, return the numeric value itself.
+    """
+    if enum is None:
+        return num
+    try:
+        return socket.AddressType(num)
+    except ValueError:
+        return num
+
+
 # --- Process.connections() 'kind' parameter mapping
 
 conn_tmap = {
@@ -184,7 +212,7 @@ if AF_UNIX is not None:
         "unix": ([AF_UNIX], [SOCK_STREAM, SOCK_DGRAM]),
     })
 
-del AF_INET, AF_INET6, AF_UNIX, SOCK_STREAM, SOCK_DGRAM, socket
+del AF_INET, AF_INET6, AF_UNIX, SOCK_STREAM, SOCK_DGRAM
 
 
 # --- namedtuples for psutil.* system-related functions
diff --git a/psutil/_psbsd.py b/psutil/_psbsd.py
index 87ae2119..9ce67441 100644
--- a/psutil/_psbsd.py
+++ b/psutil/_psbsd.py
@@ -14,7 +14,8 @@ from collections import namedtuple
 
 from psutil import _common
 from psutil import _psposix
-from psutil._common import conn_tmap, usage_percent
+from psutil._common import (conn_tmap, usage_percent, sockfam_to_enum,
+                            socktype_to_enum)
 import _psutil_bsd as cext
 import _psutil_posix
 
@@ -191,6 +192,8 @@ def net_connections(kind):
         # TODO: apply filter at C level
         if fam in families and type in types:
             status = TCP_STATUSES[status]
+            fam = sockfam_to_enum(fam)
+            type = socktype_to_enum(type)
             nt = _common.sconn(fd, fam, type, laddr, raddr, status, pid)
             ret.append(nt)
     return ret
@@ -312,6 +315,8 @@ class Process(object):
         ret = []
         for item in rawlist:
             fd, fam, type, laddr, raddr, status = item
+            fam = sockfam_to_enum(fam)
+            type = socktype_to_enum(type)
             status = TCP_STATUSES[status]
             nt = _common.pconn(fd, fam, type, laddr, raddr, status)
             ret.append(nt)
diff --git a/psutil/_psosx.py b/psutil/_psosx.py
index c40ef1d7..94cd3f61 100644
--- a/psutil/_psosx.py
+++ b/psutil/_psosx.py
@@ -13,7 +13,8 @@ from collections import namedtuple
 
 from psutil import _common
 from psutil import _psposix
-from psutil._common import conn_tmap, usage_percent, isfile_strict
+from psutil._common import (conn_tmap, usage_percent, isfile_strict,
+                            sockfam_to_enum, socktype_to_enum)
 import _psutil_osx as cext
 import _psutil_posix
 
@@ -292,6 +293,8 @@ 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)
             nt = _common.pconn(fd, fam, type, laddr, raddr, status)
             ret.append(nt)
         return ret
diff --git a/psutil/_pssunos.py b/psutil/_pssunos.py
index 7b5385ef..ab8a6357 100644
--- a/psutil/_pssunos.py
+++ b/psutil/_pssunos.py
@@ -15,7 +15,8 @@ from collections import namedtuple
 
 from psutil import _common
 from psutil import _psposix
-from psutil._common import usage_percent, isfile_strict
+from psutil._common import (usage_percent, isfile_strict, sockfam_to_enum,
+                            socktype_to_enum)
 from psutil._compat import PY3
 import _psutil_posix
 import _psutil_sunos as cext
@@ -218,6 +219,8 @@ def net_connections(kind, _pid=-1):
         if type_ not in types:
             continue
         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)
         else:
diff --git a/psutil/_pswindows.py b/psutil/_pswindows.py
index 6068446b..b1f2a45d 100644
--- a/psutil/_pswindows.py
+++ b/psutil/_pswindows.py
@@ -12,7 +12,8 @@ import os
 from collections import namedtuple
 
 from psutil import _common
-from psutil._common import conn_tmap, usage_percent, isfile_strict
+from psutil._common import (conn_tmap, usage_percent, isfile_strict,
+                            sockfam_to_enum, socktype_to_enum)
 from psutil._compat import PY3, xrange, lru_cache
 import _psutil_windows as cext
 
@@ -171,6 +172,8 @@ def net_connections(kind, _pid=-1):
     for item in rawlist:
         fd, fam, type, laddr, raddr, status, pid = item
         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)
         else:
diff --git a/test/test_psutil.py b/test/test_psutil.py
index 7979796f..9442b88c 100644
--- a/test/test_psutil.py
+++ b/test/test_psutil.py
@@ -46,6 +46,10 @@ try:
     import ipaddress  # python >= 3.3
 except ImportError:
     ipaddress = None
+try:
+    import enum  # python >= 3.4
+except ImportError:
+    enum = None
 
 if sys.version_info < (2, 7):
     import unittest2 as unittest  # https://pypi.python.org/pypi/unittest2
@@ -282,6 +286,8 @@ def reap_children(search_all=False):
 
 def check_ip_address(addr, family):
     """Attempts to check IP address's validity."""
+    if enum is not None:
+        assert isinstance(family, enum.IntEnum), family
     if family == AF_INET:
         octs = [int(x) for x in addr.split('.')]
         assert len(octs) == 4, addr