Merge branch 'master' of github.com:giampaolo/psutil

7 files changed, 49 insertions, 32 deletions

diff --git a/psutil/_psutil_common.c b/psutil/_psutil_common.c
index a4899487..6d13ec4a 100644
--- a/psutil/_psutil_common.c
+++ b/psutil/_psutil_common.c
@@ -120,10 +120,11 @@ psutil_debug(const char* format, ...) {
 /*
  * Called on module import on all platforms.
  */
-void
-psutil_setup(void) {
+int
+psutil_setup() {
     if (getenv("PSUTIL_DEBUG") != NULL)
         PSUTIL_DEBUG = 1;
     if (getenv("PSUTIL_TESTING") != NULL)
         PSUTIL_TESTING = 1;
+    return 0;
 }
diff --git a/psutil/_psutil_common.h b/psutil/_psutil_common.h
index 70698107..5e6640c2 100644
--- a/psutil/_psutil_common.h
+++ b/psutil/_psutil_common.h
@@ -26,6 +26,6 @@ PyObject* PyErr_SetFromOSErrnoWithSyscall(const char *syscall);
 
 PyObject* psutil_set_testing(PyObject *self, PyObject *args);
 void psutil_debug(const char* format, ...);
-void psutil_setup(void);
+int psutil_setup();
 
 #endif // PSUTIL_PSUTIL_COMMON_H
diff --git a/psutil/_psutil_linux.c b/psutil/_psutil_linux.c
index be808633..4bf53b85 100644
--- a/psutil/_psutil_linux.c
+++ b/psutil/_psutil_linux.c
@@ -673,6 +673,8 @@ void init_psutil_linux(void)
 #else
     PyObject *module = Py_InitModule("_psutil_linux", PsutilMethods);
 #endif
+    if (module == NULL)
+        INITERROR;
 
     PyModule_AddIntConstant(module, "version", PSUTIL_VERSION);
 #if PSUTIL_HAVE_PRLIMIT
@@ -720,8 +722,6 @@ void init_psutil_linux(void)
     PyModule_AddIntConstant(module, "DUPLEX_FULL", DUPLEX_FULL);
     PyModule_AddIntConstant(module, "DUPLEX_UNKNOWN", DUPLEX_UNKNOWN);
 
-    psutil_setup();
-
     if (module == NULL)
         INITERROR;
 #if PY_MAJOR_VERSION >= 3
diff --git a/psutil/_psutil_osx.c b/psutil/_psutil_osx.c
index 08be754b..ed4e60d3 100644
--- a/psutil/_psutil_osx.c
+++ b/psutil/_psutil_osx.c
@@ -2028,6 +2028,12 @@ init_psutil_osx(void)
 #else
     PyObject *module = Py_InitModule("_psutil_osx", PsutilMethods);
 #endif
+    if (module == NULL)
+        INITERROR;
+
+    if (psutil_setup() != 0)
+        INITERROR;
+
     PyModule_AddIntConstant(module, "version", PSUTIL_VERSION);
     // process status constants, defined in:
     // http://fxr.watson.org/fxr/source/bsd/sys/proc.h?v=xnu-792.6.70#L149
@@ -2056,8 +2062,6 @@ init_psutil_osx(void)
     Py_INCREF(ZombieProcessError);
     PyModule_AddObject(module, "ZombieProcessError", ZombieProcessError);
 
-    psutil_setup();
-
     if (module == NULL)
         INITERROR;
 #if PY_MAJOR_VERSION >= 3
diff --git a/psutil/_psutil_sunos.c b/psutil/_psutil_sunos.c
index 0717f195..99423f7a 100644
--- a/psutil/_psutil_sunos.c
+++ b/psutil/_psutil_sunos.c
@@ -1732,6 +1732,12 @@ void init_psutil_sunos(void)
 #else
     PyObject *module = Py_InitModule("_psutil_sunos", PsutilMethods);
 #endif
+    if (module == NULL)
+        INITERROR;
+
+    if (psutil_setup() != 0)
+        INITERROR;
+
     PyModule_AddIntConstant(module, "version", PSUTIL_VERSION);
 
     PyModule_AddIntConstant(module, "SSLEEP", SSLEEP);
@@ -1761,8 +1767,6 @@ void init_psutil_sunos(void)
     PyModule_AddIntConstant(module, "TCPS_BOUND", TCPS_BOUND);
     PyModule_AddIntConstant(module, "PSUTIL_CONN_NONE", PSUTIL_CONN_NONE);
 
-    psutil_setup();
-
     if (module == NULL)
         INITERROR;
 #if PY_MAJOR_VERSION >= 3
diff --git a/psutil/_psutil_windows.c b/psutil/_psutil_windows.c
index 48e75863..54b401d3 100644
--- a/psutil/_psutil_windows.c
+++ b/psutil/_psutil_windows.c
@@ -12,22 +12,15 @@
 #include <Python.h>
 #include <windows.h>
 #include <Psapi.h>
-#include <time.h>
-#include <lm.h>
-#include <WinIoCtl.h>
+#include <signal.h>
+#include <WinIoCtl.h>  // disk_io_counters()
 #include <tchar.h>
 #include <tlhelp32.h>
-#include <winsock2.h>
-#if (_WIN32_WINNT >= 0x0600) // Windows Vista and above
-#include <ws2tcpip.h>
+#include <wtsapi32.h>  // users()
+#include <PowrProf.h>  // cpu_freq()
+#if (_WIN32_WINNT >= 0x0600) // Windows >= Vista
+#include <ws2tcpip.h>  // net_connections()
 #endif
-#include <iphlpapi.h>
-#include <iprtrmib.h>
-#include <udpmib.h>
-#include <wtsapi32.h>
-#include <Winsvc.h>
-#include <PowrProf.h>
-#include <signal.h>
 
 // Link with Iphlpapi.lib
 #pragma comment(lib, "IPHLPAPI.lib")
@@ -3615,10 +3608,15 @@ void init_psutil_windows(void)
 #else
     PyObject *module = Py_InitModule("_psutil_windows", PsutilMethods);
 #endif
+    if (module == NULL)
+        INITERROR;
 
-    if (module == NULL) {
+    if (psutil_setup() != 0)
+        INITERROR;
+    if (psutil_load_globals() != 0)
+        INITERROR;
+    if (psutil_set_se_debug() != 0)
         INITERROR;
-    }
 
     st = GETSTATE(module);
     st->error = PyErr_NewException("_psutil_windows.Error", NULL, NULL);
@@ -3715,13 +3713,22 @@ void init_psutil_windows(void)
     PyModule_AddIntConstant(
         module, "ERROR_SERVICE_DOES_NOT_EXIST", ERROR_SERVICE_DOES_NOT_EXIST);
 
-    // set SeDebug for the current process
-    if (psutil_set_se_debug() != 0)
-        return NULL;
-
-    psutil_setup();
-    if (psutil_load_globals() != 0)
-        return NULL;
+    PyModule_AddIntConstant(
+        module, "WINVER", PSUTIL_WINVER);
+    PyModule_AddIntConstant(
+        module, "WINDOWS_XP", PSUTIL_WINDOWS_XP);
+    PyModule_AddIntConstant(
+        module, "WINDOWS_SERVER_2003", PSUTIL_WINDOWS_SERVER_2003);
+    PyModule_AddIntConstant(
+        module, "WINDOWS_VISTA", PSUTIL_WINDOWS_VISTA);
+    PyModule_AddIntConstant(
+        module, "WINDOWS_7", PSUTIL_WINDOWS_7);
+    PyModule_AddIntConstant(
+        module, "WINDOWS_8", PSUTIL_WINDOWS_8);
+    PyModule_AddIntConstant(
+        module, "WINDOWS_8_1", PSUTIL_WINDOWS_8_1);
+    PyModule_AddIntConstant(
+        module, "WINDOWS_10", PSUTIL_WINDOWS_10);
 
 #if PY_MAJOR_VERSION >= 3
     return module;
diff --git a/setup.py b/setup.py
index 764f8b31..465a9b9e 100755
--- a/setup.py
+++ b/setup.py
@@ -135,11 +135,12 @@ if WINDOWS:
             'psutil/arch/windows/inet_ntop.c',
             'psutil/arch/windows/services.c',
             'psutil/arch/windows/global.c',
+            # 'psutil/arch/windows/connections.c',
         ],
         define_macros=macros,
         libraries=[
             "psapi", "kernel32", "advapi32", "shell32", "netapi32",
-            "iphlpapi", "wtsapi32", "ws2_32", "PowrProf",
+            "wtsapi32", "ws2_32", "PowrProf",
         ],
         # extra_compile_args=["/Z7"],
         # extra_link_args=["/DEBUG"]