move profile_func() from fwtwirl as print_profile_func()

alphatwirl/misc/__init__.py

@@ -1,6 +1,7 @@
 from .mkdir_p import mkdir_p
 from .list_to_aligned_text import list_to_aligned_text
 from .quote_string import quote_string
+from .profile_func import profile_func, print_profile_func
 
 # to be deleted
 from .listToAlignedText import listToAlignedText

alphatwirl/misc/profile_func.py

@@ -0,0 +1,30 @@
+# Tai Sakuma <tai.sakuma@cern.ch>
+from __future__ import print_function
+import cProfile, pstats
+from io import StringIO, BytesIO
+
+##__________________________________________________________________||
+def print_profile_func(func, profile_out_path=None):
+    result = profile_func(func)
+    if profile_out_path is None:
+        print(result)
+    else:
+        with open(profile_out_path, 'w') as f:
+            f.write(result)
+
+##__________________________________________________________________||
+def profile_func(func):
+    pr = cProfile.Profile()
+    pr.enable()
+    func()
+    pr.disable()
+    sortby = 'cumulative'
+    try:
+        s = StringIO()
+        pstats.Stats(pr, stream=s).strip_dirs().sort_stats(sortby).print_stats()
+    except TypeError:
+        s = BytesIO()
+        pstats.Stats(pr, stream=s).strip_dirs().sort_stats(sortby).print_stats()
+    return s.getvalue()
+
+##__________________________________________________________________||

tests/unit/misc/test_profile_func.py

@@ -0,0 +1,27 @@
+# Tai Sakuma <tai.sakuma@gmail.com>
+from __future__ import print_function
+import os
+
+import pytest
+
+try:
+    import unittest.mock as mock
+except ImportError:
+    import mock
+
+from alphatwirl.misc import profile_func, print_profile_func
+
+##__________________________________________________________________||
+def test_profile_func():
+    func = mock.Mock()
+    profile_func(func)
+
+##__________________________________________________________________||
+def test_print_profile_func(tmpdir_factory):
+    tmpdir = str(tmpdir_factory.mktemp(''))
+    func = mock.Mock()
+    print(tmpdir)
+    profile_out_path = os.path.join(tmpdir, 'profile.txt')
+    print_profile_func(func, profile_out_path=profile_out_path)
+
+##__________________________________________________________________||