Update per code review

6 files changed, 115 insertions, 145 deletions

diff --git a/psutil/_psosx.py b/psutil/_psosx.py
index e3582b8b..52994847 100644
--- a/psutil/_psosx.py
+++ b/psutil/_psosx.py
@@ -278,7 +278,8 @@ class Process(object):
 
     @wrap_exceptions
     def memory_info_ex(self):
-        rss, vms, pfaults, pageins, uss = cext.proc_memory_info(self.pid)
+        rss, vms, pfaults, pageins = cext.proc_memory_info(self.pid)
+        uss = cext.proc_memory_uss(self.pid)
         return pextmem(rss, vms, pfaults * PAGESIZE, pageins * PAGESIZE, uss)
 
     @wrap_exceptions
diff --git a/psutil/_psutil_osx.c b/psutil/_psutil_osx.c
index ab3ac53a..0cf0684b 100644
--- a/psutil/_psutil_osx.c
+++ b/psutil/_psutil_osx.c
@@ -40,7 +40,6 @@
 #include <IOKit/IOBSD.h>
 
 #include "_psutil_osx.h"
-#include "_psutil_osx_uss.h"
 #include "_psutil_common.h"
 #include "arch/osx/process_info.h"
 
@@ -516,9 +515,6 @@ static PyObject *
 psutil_proc_memory_info(PyObject *self, PyObject *args) {
     long pid;
     struct proc_taskinfo pti;
-    int err;
-    mach_port_t task = MACH_PORT_NULL;
-    int64_t uss = 0;
 
     if (! PyArg_ParseTuple(args, "l", &pid))
         return NULL;
@@ -529,31 +525,123 @@ psutil_proc_memory_info(PyObject *self, PyObject *args) {
     // I just give up...
     // struct proc_regioninfo pri;
     // psutil_proc_pidinfo(pid, PROC_PIDREGIONINFO, &pri, sizeof(pri))
+    return Py_BuildValue(
+        "(KKkk)",
+        pti.pti_resident_size,  // resident memory size (rss)
+        pti.pti_virtual_size,   // virtual memory size (vms)
+        pti.pti_faults,         // number of page faults (pages)
+        pti.pti_pageins         // number of actual pageins (pages)
+    );
+}
 
-    err = task_for_pid(mach_task_self(), pid, &task);
+/**
+ * Indicates if the given virtual address on the given architecture is in the
+ * shared VM region.
+ */
+bool
+in_shared_region(mach_vm_address_t addr, cpu_type_t type)
+{
+    mach_vm_address_t base;
+    mach_vm_address_t size;
+
+    switch (type) {
+        case CPU_TYPE_ARM:
+            base = SHARED_REGION_BASE_ARM;
+            size = SHARED_REGION_SIZE_ARM;
+            break;
+        case CPU_TYPE_I386:
+            base = SHARED_REGION_BASE_I386;
+            size = SHARED_REGION_SIZE_I386;
+            break;
+        case CPU_TYPE_X86_64:
+            base = SHARED_REGION_BASE_X86_64;
+            size = SHARED_REGION_SIZE_X86_64;
+            break;
+        default:
+            return false;
+    }
+
+    return base <= addr && addr < (base + size);
+}
+
+
+/**
+ * Returns the USS (unique set size) of the process.
+ */
+static PyObject *
+psutil_proc_memory_uss(PyObject *self, PyObject *args)
+{
+    long pid;
+    if (! PyArg_ParseTuple(args, "l", &pid))
+        return NULL;
+
+    mach_port_t task = MACH_PORT_NULL;
+    int err = task_for_pid(mach_task_self(), pid, &task);
     if (err != KERN_SUCCESS) {
         psutil_raise_ad_or_nsp(pid);
-        goto error;
+        return NULL;
     }
 
-    calc_uss(task, &uss);
+    cpu_type_t cpu_type;
+    size_t len = sizeof(cpu_type);
+    if (sysctlbyname("sysctl.proc_cputype", &cpu_type, &len, NULL, 0) != 0) {
+        PyErr_SetFromErrno(PyExc_OSError);
+        return NULL;
+    }
 
-    if (task != MACH_PORT_NULL)
-        mach_port_deallocate(mach_task_self(), task);
+    // Roughly based on libtop_update_vm_regions in
+    // http://www.opensource.apple.com/source/top/top-100.1.2/libtop.c
+    size_t private_pages = 0;
+    mach_vm_size_t size = 0;
+    for (mach_vm_address_t addr = MACH_VM_MIN_ADDRESS; ; addr += size) {
+        vm_region_top_info_data_t info;
+        mach_msg_type_number_t info_count = VM_REGION_TOP_INFO_COUNT;
+        mach_port_t object_name;
+
+        kern_return_t kr =
+            mach_vm_region(task, &addr, &size, VM_REGION_TOP_INFO,
+                    (vm_region_info_t)&info,
+                    &info_count, &object_name);
+        if (kr == KERN_INVALID_ADDRESS) {
+            // Done iterating VM regions.
+            break;
+        }
+        else if (kr != KERN_SUCCESS) {
+            return false;
+        }
 
-    return Py_BuildValue(
-        "(KKkkK)",
-        pti.pti_resident_size,  // resident memory size (rss)
-        pti.pti_virtual_size,   // virtual memory size (vms)
-        pti.pti_faults,         // number of page faults (pages)
-        pti.pti_pageins,        // number of actual pageins (pages)
-        (long long)uss          // unique memory size (uss)
-    );
+        if (in_shared_region(addr, cpu_type) && info.share_mode != SM_PRIVATE) {
+            continue;
+        }
 
-error:
-    if (task != MACH_PORT_NULL)
-        mach_port_deallocate(mach_task_self(), task);
-    return NULL;
+        switch (info.share_mode) {
+            case SM_LARGE_PAGE:
+                // NB: Large pages are not shareable and always resident.
+            case SM_PRIVATE:
+                private_pages += info.private_pages_resident;
+                private_pages += info.shared_pages_resident;
+                break;
+            case SM_COW:
+                private_pages += info.private_pages_resident;
+                if (info.ref_count == 1) {
+                    // Treat copy-on-write pages as private if they only have one reference.
+                    private_pages += info.shared_pages_resident;
+                }
+                break;
+            case SM_SHARED:
+            default:
+                break;
+        }
+    }
+
+    mach_port_deallocate(mach_task_self(), task);
+
+    vm_size_t page_size;
+    if (host_page_size(mach_host_self(), &page_size) != KERN_SUCCESS) {
+        page_size = PAGE_SIZE;
+    }
+
+    return Py_BuildValue("K", private_pages * page_size);
 }
 
 
@@ -1722,6 +1810,8 @@ PsutilMethods[] = {
      "seconds since the epoch"},
     {"proc_memory_info", psutil_proc_memory_info, METH_VARARGS,
      "Return memory information about a process"},
+    {"proc_memory_uss", psutil_proc_memory_uss, METH_VARARGS,
+     "Return USS a process"},
     {"proc_num_threads", psutil_proc_num_threads, METH_VARARGS,
      "Return number of threads used by process"},
     {"proc_status", psutil_proc_status, METH_VARARGS,
diff --git a/psutil/_psutil_osx.h b/psutil/_psutil_osx.h
index 2220ec3b..a6a2957f 100644
--- a/psutil/_psutil_osx.h
+++ b/psutil/_psutil_osx.h
@@ -17,6 +17,7 @@ static PyObject* psutil_proc_exe(PyObject* self, PyObject* args);
 static PyObject* psutil_proc_gids(PyObject* self, PyObject* args);
 static PyObject* psutil_proc_memory_info(PyObject* self, PyObject* args);
 static PyObject* psutil_proc_memory_maps(PyObject* self, PyObject* args);
+static PyObject* psutil_proc_memory_uss(PyObject* self, PyObject* args);
 static PyObject* psutil_proc_name(PyObject* self, PyObject* args);
 static PyObject* psutil_proc_num_fds(PyObject* self, PyObject* args);
 static PyObject* psutil_proc_num_threads(PyObject* self, PyObject* args);
diff --git a/psutil/_psutil_osx_uss.c b/psutil/_psutil_osx_uss.c
deleted file mode 100644
index 72912f52..00000000
--- a/psutil/_psutil_osx_uss.c
+++ /dev/null
@@ -1,105 +0,0 @@
-/* -*- Mode: C++; tab-width: 8; indent-tabs-mode: nil; c-basic-offset: 2 -*- */
-/* vim: set ts=8 sts=2 et sw=2 tw=80: */
-/* This Source Code Form is subject to the terms of the Mozilla Public
- * License, v. 2.0. If a copy of the MPL was not distributed with this
- * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
-
-#include <mach/mach_init.h>
-#include <mach/mach_vm.h>
-#include <mach/shared_region.h>
-#include <mach/task.h>
-#include <sys/sysctl.h>
-
-#include "_psutil_osx_uss.h"
-
-bool
-InSharedRegion(mach_vm_address_t aAddr, cpu_type_t aType)
-{
-  mach_vm_address_t base;
-  mach_vm_address_t size;
-
-  switch (aType) {
-    case CPU_TYPE_ARM:
-      base = SHARED_REGION_BASE_ARM;
-      size = SHARED_REGION_SIZE_ARM;
-      break;
-    case CPU_TYPE_I386:
-      base = SHARED_REGION_BASE_I386;
-      size = SHARED_REGION_SIZE_I386;
-      break;
-    case CPU_TYPE_X86_64:
-      base = SHARED_REGION_BASE_X86_64;
-      size = SHARED_REGION_SIZE_X86_64;
-      break;
-    default:
-      return false;
-  }
-
-  return base <= aAddr && aAddr < (base + size);
-}
-
-bool
-calc_uss(mach_port_t target, int64_t* aN)
-{
-  if (!aN) {
-    return false;
-  }
-
-  cpu_type_t cpu_type;
-  size_t len = sizeof(cpu_type);
-  if (sysctlbyname("sysctl.proc_cputype", &cpu_type, &len, NULL, 0) != 0) {
-    return false;
-  }
-
-  /* Roughly based on libtop_update_vm_regions in
-    http://www.opensource.apple.com/source/top/top-100.1.2/libtop.c */
-  size_t privatePages = 0;
-  mach_vm_size_t size = 0;
-  for (mach_vm_address_t addr = MACH_VM_MIN_ADDRESS; ; addr += size) {
-    vm_region_top_info_data_t info;
-    mach_msg_type_number_t infoCount = VM_REGION_TOP_INFO_COUNT;
-    mach_port_t objectName;
-
-    kern_return_t kr =
-        mach_vm_region(target, &addr, &size, VM_REGION_TOP_INFO,
-                       (vm_region_info_t)&info,
-                       &infoCount, &objectName);
-    if (kr == KERN_INVALID_ADDRESS) {
-      /* Done iterating VM regions. */
-      break;
-    } else if (kr != KERN_SUCCESS) {
-      return false;
-    }
-
-    if (InSharedRegion(addr, cpu_type) && info.share_mode != SM_PRIVATE) {
-        continue;
-    }
-
-    switch (info.share_mode) {
-      case SM_LARGE_PAGE:
-        /* NB: Large pages are not shareable and always resident. */
-      case SM_PRIVATE:
-        privatePages += info.private_pages_resident;
-        privatePages += info.shared_pages_resident;
-        break;
-      case SM_COW:
-        privatePages += info.private_pages_resident;
-        if (info.ref_count == 1) {
-          /* Treat copy-on-write pages as private if they only have one reference. */
-          privatePages += info.shared_pages_resident;
-        }
-        break;
-      case SM_SHARED:
-      default:
-        break;
-    }
-  }
-
-  vm_size_t pageSize;
-  if (host_page_size(mach_host_self(), &pageSize) != KERN_SUCCESS) {
-    pageSize = PAGE_SIZE;
-  }
-
-  *aN = privatePages * pageSize;
-  return true;
-}
diff --git a/psutil/_psutil_osx_uss.h b/psutil/_psutil_osx_uss.h
deleted file mode 100644
index 87e75f6f..00000000
--- a/psutil/_psutil_osx_uss.h
+++ /dev/null
@@ -1,16 +0,0 @@
-/* -*- Mode: C++; tab-width: 8; indent-tabs-mode: nil; c-basic-offset: 2 -*- */
-/* vim: set ts=8 sts=2 et sw=2 tw=80: */
-/* This Source Code Form is subject to the terms of the Mozilla Public
- * License, v. 2.0. If a copy of the MPL was not distributed with this
- * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
-
-#ifndef _psutil_osx_uss_h
-#define _psutil_osx_uss_h
-
-#include <stdbool.h>
-#include <stdint.h>
-#include <mach/port.h>
-
-bool calc_uss(mach_port_t target, int64_t* aN);
-
-#endif
diff --git a/setup.py b/setup.py
index aef6abdd..4d4d8650 100644
--- a/setup.py
+++ b/setup.py
@@ -119,8 +119,7 @@ elif sys.platform.startswith("darwin"):
         sources=[
             'psutil/_psutil_osx.c',
             'psutil/_psutil_common.c',
-            'psutil/arch/osx/process_info.c',
-            'psutil/_psutil_osx_uss.c'
+            'psutil/arch/osx/process_info.c'
         ],
         define_macros=[VERSION_MACRO],
         extra_link_args=[