refactor import

4 files changed, 100 insertions, 94 deletions

diff --git a/docs/index.rst b/docs/index.rst
index 3898f9da..58df4e99 100644
--- a/docs/index.rst
+++ b/docs/index.rst
@@ -1348,12 +1348,9 @@ Process class
 
       >>> import psutil
       >>> p = psutil.Process()
-      >>> # process may open no more than 128 file descriptors
-      >>> p.rlimit(psutil.RLIMIT_NOFILE, (128, 128))
-      >>> # process may create files no bigger than 1024 bytes
-      >>> p.rlimit(psutil.RLIMIT_FSIZE, (1024, 1024))
-      >>> # get
-      >>> p.rlimit(psutil.RLIMIT_FSIZE)
+      >>> p.rlimit(psutil.RLIMIT_NOFILE, (128, 128))   # process can open max 128 file descriptors
+      >>> p.rlimit(psutil.RLIMIT_FSIZE, (1024, 1024))  # can create files no bigger than 1024 bytes
+      >>> p.rlimit(psutil.RLIMIT_FSIZE)                # get
       (1024, 1024)
       >>>
 
@@ -2253,23 +2250,35 @@ Process resources constants
 .. data:: RLIMIT_CPU
 .. data:: RLIMIT_DATA
 .. data:: RLIMIT_FSIZE
-.. data:: RLIMIT_LOCKS
 .. data:: RLIMIT_MEMLOCK
-.. data:: RLIMIT_MSGQUEUE
-.. data:: RLIMIT_NICE
 .. data:: RLIMIT_NOFILE
 .. data:: RLIMIT_NPROC
 .. data:: RLIMIT_RSS
-.. data:: RLIMIT_RTPRIO
-.. data:: RLIMIT_RTTIME
-.. data:: RLIMIT_SIGPENDING
 .. data:: RLIMIT_STACK
 
-  Constants used for getting and setting process resource limits to be used in
-  conjunction with :meth:`psutil.Process.rlimit()`. See `man prlimit`_ for
-  further information.
+Linux specific:
 
-  Availability: Linux
+  .. data:: RLIMIT_LOCKS
+  .. data:: RLIMIT_MSGQUEUE
+  .. data:: RLIMIT_NICE
+  .. data:: RLIMIT_RTPRIO
+  .. data:: RLIMIT_RTTIME
+  .. data:: RLIMIT_SIGPENDING
+
+FreeBSD specific:
+
+  .. data:: RLIMIT_SWAP
+  .. data:: RLIMIT_SBSIZE
+  .. data:: RLIMIT_NPTS
+
+Constants used for getting and setting process resource limits to be used in
+conjunction with :meth:`psutil.Process.rlimit()`. See `resource.getrlimit`_
+for further information.
+
+Availability: Linux, FreeBSD
+
+.. versionchanged:: 5.7.3 added FreeBSD support, added ``RLIMIT_SWAP``,
+  ``RLIMIT_SBSIZE``, ``RLIMIT_NPTS``.
 
 Connections constants
 ---------------------
diff --git a/psutil/__init__.py b/psutil/__init__.py
index 9951692b..34e7649a 100644
--- a/psutil/__init__.py
+++ b/psutil/__init__.py
@@ -145,78 +145,54 @@ else:  # pragma: no cover
 
 if LINUX or FREEBSD:
     # Kinda ugly but considerably faster than using hasattr() and
-    # setattr() against the module object (we are at import time:
-    # speed matters).
+    # setattr() against the module object or doing it dynamically
+    # via __import__ (we are at import time: speed matters).
     from . import _psutil_posix
 
-    try:
-        RLIM_INFINITY = _psutil_posix.RLIM_INFINITY  # NOQA
-    except AttributeError:
-        pass
-    try:
-        RLIMIT_AS = _psutil_posix.RLIMIT_AS  # NOQA
-    except AttributeError:
-        pass
-    try:
-        RLIMIT_CORE = _psutil_posix.RLIMIT_CORE  # NOQA
-    except AttributeError:
-        pass
-    try:
-        RLIMIT_CPU = _psutil_posix.RLIMIT_CPU  # NOQA
-    except AttributeError:
-        pass
-    try:
-        RLIMIT_DATA = _psutil_posix.RLIMIT_DATA  # NOQA
-    except AttributeError:
-        pass
-    try:
-        RLIMIT_FSIZE = _psutil_posix.RLIMIT_FSIZE  # NOQA
-    except AttributeError:
-        pass
-    try:
-        RLIMIT_LOCKS = _psutil_posix.RLIMIT_LOCKS  # NOQA
-    except AttributeError:
-        pass
-    try:
-        RLIMIT_MEMLOCK = _psutil_posix.RLIMIT_MEMLOCK  # NOQA
-    except AttributeError:
-        pass
-    try:
-        RLIMIT_NOFILE = _psutil_posix.RLIMIT_NOFILE  # NOQA
-    except AttributeError:
-        pass
-    try:
-        RLIMIT_NPROC = _psutil_posix.RLIMIT_NPROC  # NOQA
-    except AttributeError:
-        pass
-    try:
-        RLIMIT_RSS = _psutil_posix.RLIMIT_RSS  # NOQA
-    except AttributeError:
-        pass
-    try:
-        RLIMIT_STACK = _psutil_posix.RLIMIT_STACK  # NOQA
-    except AttributeError:
-        pass
-    try:
-        RLIMIT_MSGQUEUE = _psutil_posix.RLIMIT_MSGQUEUE
-    except AttributeError:
-        pass
-    try:
-        RLIMIT_NICE = _psutil_posix.RLIMIT_NICE
-    except AttributeError:
-        pass
-    try:
-        RLIMIT_RTPRIO = _psutil_posix.RLIMIT_RTPRIO
-    except AttributeError:
-        pass
-    try:
-        RLIMIT_RTTIME = _psutil_posix.RLIMIT_RTTIME
-    except AttributeError:
-        pass
-    try:
-        RLIMIT_SIGPENDING = _psutil_posix.RLIMIT_SIGPENDING
-    except AttributeError:
-        pass
+    _dir = set(_psplatform.__extra__all__)
+    if 'RLIM_INFINITY' in _dir:
+        RLIM_INFINITY = _psutil_posix.RLIM_INFINITY
+    if 'RLIMIT_AS' in _dir:
+        RLIMIT_AS = _psutil_posix.RLIMIT_AS
+    if 'RLIMIT_CORE' in _dir:
+        RLIMIT_CORE = _psutil_posix.RLIMIT_CORE
+    if 'RLIMIT_CPU' in _dir:
+        RLIMIT_CPU = _psutil_posix.RLIMIT_CPU
+    if 'RLIMIT_DATA' in _dir:
+        RLIMIT_DATA = _psutil_posix.RLIMIT_DATA
+    if 'RLIMIT_FSIZE' in _dir:
+        RLIMIT_FSIZE = _psutil_posix.RLIMIT_FSIZE
+    if 'RLIMIT_MEMLOCK' in _dir:
+        RLIMIT_MEMLOCK = _psutil_posix.RLIMIT_MEMLOCK
+    if 'RLIMIT_NOFILE' in _dir:
+        RLIMIT_NOFILE = _psutil_posix.RLIMIT_NOFILE
+    if 'RLIMIT_NPROC' in _dir:
+        RLIMIT_NPROC = _psutil_posix.RLIMIT_NPROC
+    if 'RLIMIT_RSS' in _dir:
+        RLIMIT_RSS = _psutil_posix.RLIMIT_RSS
+    if 'RLIMIT_STACK' in _dir:
+        RLIMIT_STACK = _psutil_posix.RLIMIT_STACK
+    if LINUX:
+        if 'RLIMIT_LOCKS' in _dir:
+            RLIMIT_LOCKS = _psutil_posix.RLIMIT_LOCKS
+        if 'RLIMIT_MSGQUEUE' in _dir:
+            RLIMIT_MSGQUEUE = _psutil_posix.RLIMIT_MSGQUEUE
+        if 'RLIMIT_NICE' in _dir:
+            RLIMIT_NICE = _psutil_posix.RLIMIT_NICE
+        if 'RLIMIT_RTPRIO' in _dir:
+            RLIMIT_RTPRIO = _psutil_posix.RLIMIT_RTPRIO
+        if 'RLIMIT_RTTIME' in _dir:
+            RLIMIT_RTTIME = _psutil_posix.RLIMIT_RTTIME
+        if 'RLIMIT_SIGPENDING' in _dir:
+            RLIMIT_SIGPENDING = _psutil_posix.RLIMIT_SIGPENDING
+    elif FREEBSD:
+        if 'RLIMIT_SWAP' in _dir:
+            RLIMIT_SWAP = _psutil_posix.RLIMIT_SWAP
+        if 'RLIMIT_SBSIZE' in _dir:
+            RLIMIT_SBSIZE = _psutil_posix.RLIMIT_SBSIZE
+        if 'RLIMIT_NPTS' in _dir:
+            RLIMIT_NPTS = _psutil_posix.RLIMIT_NPTS
+    del _dir
 
 
 __all__ = [
@@ -235,6 +211,7 @@ __all__ = [
     "CONN_ESTABLISHED", "CONN_SYN_SENT", "CONN_SYN_RECV", "CONN_FIN_WAIT1",
     "CONN_FIN_WAIT2", "CONN_TIME_WAIT", "CONN_CLOSE", "CONN_CLOSE_WAIT",
     "CONN_LAST_ACK", "CONN_LISTEN", "CONN_CLOSING", "CONN_NONE",
+    # "CONN_IDLE", "CONN_BOUND",
 
     "AF_LINK",
 
@@ -245,6 +222,11 @@ __all__ = [
     "BSD", "FREEBSD", "LINUX", "NETBSD", "OPENBSD", "MACOS", "OSX", "POSIX",
     "SUNOS", "WINDOWS", "AIX",
 
+    # "RLIM_INFINITY", "RLIMIT_AS", "RLIMIT_CORE", "RLIMIT_CPU", "RLIMIT_DATA",
+    # "RLIMIT_FSIZE", "RLIMIT_LOCKS", "RLIMIT_MEMLOCK", "RLIMIT_NOFILE",
+    # "RLIMIT_NPROC", "RLIMIT_RSS", "RLIMIT_STACK", "RLIMIT_MSGQUEUE",
+    # "RLIMIT_NICE", "RLIMIT_RTPRIO", "RLIMIT_RTTIME", "RLIMIT_SIGPENDING",
+
     # classes
     "Process", "Popen",
 
@@ -839,7 +821,7 @@ class Process(object):
             else:
                 return self._proc.ionice_set(ioclass, value)
 
-    # Linux only
+    # Linux / FreeBSD only
     if hasattr(_psplatform.Process, "rlimit"):
 
         def rlimit(self, resource, limits=None):
@@ -847,15 +829,12 @@ class Process(object):
             tuple.
 
             *resource* is one of the RLIMIT_* constants.
-            *limits* is supposed to be a (soft, hard)  tuple.
+            *limits* is supposed to be a (soft, hard) tuple.
 
             See "man prlimit" for further info.
-            Available on Linux only.
+            Available on Linux and FreeBSD only.
             """
-            if limits is None:
-                return self._proc.rlimit(resource)
-            else:
-                return self._proc.rlimit(resource, limits)
+            return self._proc.rlimit(resource, limits)
 
     # Windows, Linux and FreeBSD only
     if hasattr(_psplatform.Process, "cpu_affinity_get"):
@@ -869,7 +848,7 @@ class Process(object):
             (Windows, Linux and BSD only).
             """
             if cpus is None:
-                return list(set(self._proc.cpu_affinity_get()))
+                return sorted(set(self._proc.cpu_affinity_get()))
             else:
                 if not cpus:
                     if hasattr(self._proc, "_get_eligible_cpus"):
diff --git a/psutil/_psutil_posix.c b/psutil/_psutil_posix.c
index d5567d3c..f1d2afd6 100644
--- a/psutil/_psutil_posix.c
+++ b/psutil/_psutil_posix.c
@@ -738,6 +738,20 @@ static PyMethodDef mod_methods[] = {
     if (PyModule_AddIntConstant(mod, "RLIMIT_SIGPENDING", RLIMIT_SIGPENDING)) INITERR;
 #endif
 
+// Free specific
+
+#ifdef RLIMIT_SWAP
+    if (PyModule_AddIntConstant(mod, "RLIMIT_SWAP", RLIMIT_SWAP)) INITERR;
+#endif
+
+#ifdef RLIMIT_SBSIZE
+    if (PyModule_AddIntConstant(mod, "RLIMIT_SBSIZE", RLIMIT_SBSIZE)) INITERR;
+#endif
+
+#ifdef RLIMIT_NPTS
+    if (PyModule_AddIntConstant(mod, "RLIMIT_NPTS", RLIMIT_NPTS)) INITERR;
+#endif
+
 #if defined(HAVE_LONG_LONG)
     if (sizeof(RLIM_INFINITY) > sizeof(long)) {
         v = PyLong_FromLongLong((PY_LONG_LONG) RLIM_INFINITY);
diff --git a/psutil/tests/test_contracts.py b/psutil/tests/test_contracts.py
index b72bcd7b..39a52569 100755
--- a/psutil/tests/test_contracts.py
+++ b/psutil/tests/test_contracts.py
@@ -108,6 +108,10 @@ class TestAvailConstantsAPIs(PsutilTestCase):
         ae(hasattr(psutil, "RLIMIT_RTTIME"), LINUX)  # requires Linux 2.6.25
         ae(hasattr(psutil, "RLIMIT_SIGPENDING"), LINUX)  # requires Linux 2.6.8
 
+        ae(hasattr(psutil, "RLIMIT_SWAP"), FREEBSD)
+        ae(hasattr(psutil, "RLIMIT_SBSIZE"), FREEBSD)
+        ae(hasattr(psutil, "RLIMIT_NPTS"), FREEBSD)
+
 
 class TestAvailSystemAPIs(PsutilTestCase):