commit eaba27d496609a4fdd9b8270793f621d219ad323 Author: Damian Johnson <atagar@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))