[tor-commits] [stem/master] Add a size_of function to get memory usage

[atagar at torproject.org]

commit eaba27d496609a4fdd9b8270793f621d219ad323
Author: Damian Johnson <atagar at torproject.org>
Date:   Mon Aug 21 12:08:05 2017 -0700

    Add a size_of function to get memory usage
    
    Presently investigating nyx's memory usage but turns out sys.getsizeof doesn't
    understand how to recurse collections. Adding a simple helper to do so.
---
 docs/change_log.rst      |  1 +
 stem/util/system.py      | 49 ++++++++++++++++++++++++++++++++++++++++++++++++
 test/unit/util/system.py |  9 +++++++++
 3 files changed, 59 insertions(+)

diff --git a/docs/change_log.rst b/docs/change_log.rst
index 59b546ad..204d5718 100644
--- a/docs/change_log.rst
+++ b/docs/change_log.rst
@@ -80,6 +80,7 @@ The following are only available within Stem's `git repository
   * Supporing pid arguments in :func:`~stem.util.system.is_running`
   * Normalized :func:`~stem.util.term.format` to return unicode
   * Don't load vim swap files as configurations
+  * Added :func:`~stem.util.system.size_of`
 
  * **Interpreter**
 
diff --git a/stem/util/system.py b/stem/util/system.py
index b3b57da9..0569992b 100644
--- a/stem/util/system.py
+++ b/stem/util/system.py
@@ -25,6 +25,7 @@ best-effort, providing **None** if the lookup fails.
 
   is_available - determines if a command is available on this system
   is_running - determines if a given process is running
+  size_of - provides the memory usage of an object
   call - runs the given system command and provides back the results
 
   name_by_pid - gets the name for a process by the given pid
@@ -62,15 +63,18 @@ best-effort, providing **None** if the lookup fails.
   ====================  ===========
 """
 
+import collections
 import ctypes
 import ctypes.util
 import distutils.spawn
+import itertools
 import mimetypes
 import multiprocessing
 import os
 import platform
 import re
 import subprocess
+import sys
 import tarfile
 import threading
 import time
@@ -89,6 +93,17 @@ State = stem.util.enum.UppercaseEnum(
   'FAILED',
 )
 
+DEFAULT_SIZE = sys.getsizeof(0)  # estimate if object lacks a __sizeof__
+
+SIZE_RECURSES = {
+  tuple: iter,
+  list: iter,
+  collections.deque: iter,
+  dict: lambda d: itertools.chain.from_iterable(d.items()),
+  set: iter,
+  frozenset: iter,
+}
+
 # Mapping of commands to if they're available or not.
 
 CMD_AVAILABLE_CACHE = {}
@@ -422,6 +437,40 @@ def is_running(command):
   return None
 
 
+def size_of(obj, exclude = None):
+  """
+  Provides the `approximate memory usage of an object
+  <https://code.activestate.com/recipes/577504/>`_. This can recurse tuples,
+  lists, deques, dicts, and sets. To teach this function to inspect additional
+  object types expand SIZE_RECURSES...
+
+  ::
+
+    stem.util.system.SIZE_RECURSES[SomeClass] = SomeClass.get_elements
+
+  .. versionadded:: 1.6.0
+
+  :param object obj: object to provide the size of
+  :param set exclude: object ids to exclude from size estimation
+
+  :returns: **int** with the size of the object in bytes
+  """
+
+  if exclude is None:
+    exclude = set()
+  elif id(obj) in exclude:
+    return 0
+
+  size = sys.getsizeof(obj, DEFAULT_SIZE)
+  exclude.add(id(obj))
+
+  if type(obj) in SIZE_RECURSES:
+    for entry in SIZE_RECURSES[type(obj)](obj):
+      size += size_of(obj, exclude)
+
+  return size
+
+
 def name_by_pid(pid):
   """
   Attempts to determine the name a given process is running under (not
diff --git a/test/unit/util/system.py b/test/unit/util/system.py
index e519aec8..a331aab7 100644
--- a/test/unit/util/system.py
+++ b/test/unit/util/system.py
@@ -148,6 +148,15 @@ class TestSystem(unittest.TestCase):
     self.assertFalse(system.is_running('irssi'))
     self.assertEqual(None, system.is_running('irssi'))
 
+  def test_size_of(self):
+    """
+    Exercises the size_of function.
+    """
+
+    self.assertTrue(10 < system.size_of('hello') < 50)
+    self.assertTrue(10 < system.size_of([]) < 50)
+    self.assertTrue(system.size_of([]) < system.size_of(['hello']) < system.size_of(['hello', 'world']))
+
   @patch('stem.util.system.call')
   @patch('stem.util.proc.is_available', Mock(return_value = False))
   @patch('stem.util.system.is_available', Mock(return_value = True))