diff --git a/.gitignore b/.gitignore index b0ec204..9c42d3f 100644 --- a/.gitignore +++ b/.gitignore @@ -11,3 +11,4 @@ mprofile_*.dat # virtual environment venv/ +.python-version diff --git a/.travis.yml b/.travis.yml index cd260b0..11d4ac0 100644 --- a/.travis.yml +++ b/.travis.yml @@ -1,6 +1,5 @@ language: python python: - - "2.7" - "3.4" - "3.5" - "3.6" diff --git a/Makefile b/Makefile index e95a767..59d4950 100644 --- a/Makefile +++ b/Makefile @@ -1,6 +1,6 @@ PYTHON ?= python -.PHONY: test +.PHONY: test develop test: $(PYTHON) -m memory_profiler test/test_func.py @@ -18,6 +18,7 @@ test: $(PYTHON) test/test_exception.py $(PYTHON) test/test_exit_code.py $(PYTHON) test/test_mprof.py + $(PYTHON) test/test_async.py develop: pip install -e . diff --git a/memory_profiler.py b/memory_profiler.py index 632bee3..e13786b 100644 --- a/memory_profiler.py +++ b/memory_profiler.py @@ -7,7 +7,10 @@ _CMD_USAGE = "python -m memory_profiler script_file.py" -from functools import wraps +from asyncio import coroutine, iscoroutinefunction +from contextlib import contextmanager +from functools import partial, wraps +import builtins import inspect import linecache import logging @@ -18,7 +21,6 @@ import time import traceback import warnings -import contextlib if sys.platform == "win32": # any value except signal.CTRL_C_EVENT and signal.CTRL_BREAK_EVENT @@ -43,17 +45,8 @@ line_cell_magic = lambda func: func magics_class = lambda cls: cls -PY2 = sys.version_info[0] == 2 - _TWO_20 = float(2 ** 20) -if PY2: - import __builtin__ as builtins - to_str = lambda x: x - from future_builtins import filter -else: - import builtins - to_str = lambda x: str(x) # .. get available packages .. try: @@ -580,7 +573,7 @@ def f(*args, **kwds): return f - @contextlib.contextmanager + @contextmanager def call_on_stack(self, func, *args, **kwds): self.current_stack_level += 1 self.stack[func].append(self.current_stack_level) @@ -701,16 +694,27 @@ def add_function(self, func): else: self.code_map.add(code) + @contextmanager + def _count_ctxmgr(self): + self.enable_by_count() + try: + yield + finally: + self.disable_by_count() + def wrap_function(self, func): """ Wrap a function to profile it. """ - def f(*args, **kwds): - self.enable_by_count() - try: - return func(*args, **kwds) - finally: - self.disable_by_count() + if iscoroutinefunction(func): + @coroutine + def f(*args, **kwargs): + with self._count_ctxmgr(): + yield from func(*args, **kwargs) + else: + def f(*args, **kwds): + with self._count_ctxmgr(): + return func(*args, **kwds) return f @@ -831,7 +835,7 @@ def show_results(prof, stream=None, precision=1): inc = u'' occurences = u'' tmp = template.format(lineno, total_mem, inc, occurences, all_lines[lineno - 1]) - stream.write(to_str(tmp)) + stream.write(tmp) stream.write(u'\n\n') @@ -1121,12 +1125,23 @@ def profile(func=None, stream=None, precision=1, backend='psutil'): if not tracemalloc.is_tracing(): tracemalloc.start() if func is not None: - @wraps(func) - def wrapper(*args, **kwargs): - prof = LineProfiler(backend=backend) - val = prof(func)(*args, **kwargs) - show_results(prof, stream=stream, precision=precision) - return val + get_prof = partial(LineProfiler, backend=backend) + show_results_bound = partial( + show_results, stream=stream, precision=precision + ) + if iscoroutinefunction(func): + @coroutine + def wrapper(*args, **kwargs): + prof = get_prof() + val = yield from prof(func)(*args, **kwargs) + show_results_bound(prof) + return val + else: + def wrapper(*args, **kwargs): + prof = get_prof() + val = prof(func)(*args, **kwargs) + show_results_bound(prof) + return val return wrapper else: @@ -1198,16 +1213,13 @@ def run_module_with_profiler(module, profiler, backend, passed_args=[]): ns = dict(_CLEAN_GLOBALS, profile=profiler) _backend = choose_backend(backend) sys.argv = [module] + passed_args - if PY2: + if _backend == 'tracemalloc' and has_tracemalloc: + tracemalloc.start() + try: run_module(module, run_name="__main__", init_globals=ns) - else: - if _backend == 'tracemalloc' and has_tracemalloc: - tracemalloc.start() - try: - run_module(module, run_name="__main__", init_globals=ns) - finally: - if has_tracemalloc and tracemalloc.is_tracing(): - tracemalloc.stop() + finally: + if has_tracemalloc and tracemalloc.is_tracing(): + tracemalloc.stop() class LogFile(object): diff --git a/setup.py b/setup.py index 3a89ee6..d53db18 100644 --- a/setup.py +++ b/setup.py @@ -1,22 +1,14 @@ -import os -import io import re from setuptools import setup # https://packaging.python.org/guides/single-sourcing-package-version/ -def read(*names, **kwargs): - with io.open( - os.path.join(os.path.dirname(__file__), *names), - encoding=kwargs.get("encoding", "utf8") - ) as fp: - return fp.read() - - -def find_version(*file_paths): - version_file = read(*file_paths) - version_match = re.search(r"^__version__ = ['\"]([^'\"]*)['\"]", - version_file, re.M) +def find_version(file_paths): + with open(file_paths) as f: + version_file = f.read() + version_match = re.search( + r"^__version__ = ['\"]([^'\"]*)['\"]", version_file, re.M + ) if version_match: return version_match.group(1) @@ -27,18 +19,14 @@ def find_version(*file_paths): Intended Audience :: Developers License :: OSI Approved :: BSD License Programming Language :: Python -Programming Language :: Python :: 2 -Programming Language :: Python :: 2.6 -Programming Language :: Python :: 2.7 Programming Language :: Python :: 3 -Programming Language :: Python :: 3.2 -Programming Language :: Python :: 3.3 Topic :: Software Development Operating System :: POSIX Operating System :: Unix """ + setup( name='memory_profiler', description='A module for monitoring memory usage of a python program', @@ -49,9 +37,10 @@ def find_version(*file_paths): url='https://github.com/pythonprofilers/memory_profiler', py_modules=['memory_profiler', 'mprof'], entry_points={ - 'console_scripts' : ['mprof = mprof:main'], + 'console_scripts': ['mprof = mprof:main'], }, install_requires=['psutil'], + python_requires='>=3.4', classifiers=[_f for _f in CLASSIFIERS.split('\n') if _f], license='BSD' ) diff --git a/test/test_async.py b/test/test_async.py new file mode 100644 index 0000000..bfafa79 --- /dev/null +++ b/test/test_async.py @@ -0,0 +1,17 @@ +import asyncio + +from memory_profiler import profile + + +@profile +@asyncio.coroutine +def my_func(): + a = [1] * (10 ** 6) + b = [2] * (2 * 10 ** 7) + yield from asyncio.sleep(1e-2) + del b + + +if __name__ == '__main__': + loop = asyncio.get_event_loop() + loop.run_until_complete(my_func())