Repository: pcdshub/pcdsutils

Issue:
Profiler should include class methods as well
`isfunction` appears to be insufficient:
https://github.com/pcdshub/pcdsutils/blob/69cf8f7ad027e13ef56c5ebe0098aec4b0e93597/pcdsutils/profile.py#L416

`inspect.ismethod` should be added as well, per the prototype of this functionality in typhos https://github.com/pcdshub/typhos/pull/505

Comments:

Patch:
diff --git a/pcdsutils/profile.py b/pcdsutils/profile.py
index 7916aaf..d2e80e2 100644
--- a/pcdsutils/profile.py
+++ b/pcdsutils/profile.py
@@ -8,9 +8,10 @@ import logging
 import pkgutil
 import warnings
 from contextlib import contextmanager
-from inspect import isclass, isfunction
+from inspect import getmembers, isclass, isfunction, ismethod
 from types import ModuleType
-from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple
+from typing import (Any, Callable, Dict, Iterable, Iterator, List, Optional,
+                    Tuple)
 
 logger = logging.getLogger(__name__)
 
@@ -54,7 +55,7 @@ def profiler_context(
     use_global_profiler: bool = False,
     output_now: bool = True,
     min_threshold: float = 0,
-) -> LineProfiler:
+) -> Iterator[LineProfiler]:
     """
     Context manager for profiling a fixed span of an application.
 
@@ -397,7 +398,7 @@ def get_native_methods(
     native_methods = set()
     if seen is None:
         seen = set()
-    for obj in module_or_cls.__dict__.values():
+    for _, obj in getmembers(module_or_cls):
         try:
             if obj in seen:
                 continue
@@ -413,7 +414,7 @@ def get_native_methods(
         if isclass(obj):
             inner_methods = get_native_methods(obj, module, seen=seen)
             native_methods.update(inner_methods)
-        elif isfunction(obj):
+        elif isfunction(obj) or ismethod(obj):
             native_methods.add(obj)
     return native_methods
 


Test patch:
diff --git a/pcdsutils/tests/dummy_submodule.py b/pcdsutils/tests/dummy_submodule.py
new file mode 100644
index 0000000..e5783ec
--- /dev/null
+++ b/pcdsutils/tests/dummy_submodule.py
@@ -0,0 +1,24 @@
+"""
+Submodule to import from for test_profile
+"""
+from inspect import getmembers
+
+
+class SomeClass:
+    def __init__(self):
+        self.some_attr = 1
+
+    def method(self):
+        return getmembers(self)
+
+    @classmethod
+    def cls_method(cls):
+        return 2
+
+    @staticmethod
+    def stat_method():
+        return 3
+
+
+def some_function():
+    return 4
diff --git a/pcdsutils/tests/test_profile.py b/pcdsutils/tests/test_profile.py
new file mode 100644
index 0000000..f4c7b81
--- /dev/null
+++ b/pcdsutils/tests/test_profile.py
@@ -0,0 +1,66 @@
+"""
+Check some of the profiler internal utilities.
+"""
+import logging
+import os.path
+
+import pytest
+
+from ..profile import (get_native_functions, get_submodules, is_native,
+                       profiler_context)
+from . import dummy_submodule
+
+logger = logging.getLogger(__name__)
+
+
+def test_is_native():
+    # Was defined there
+    assert is_native(dummy_submodule.some_function, dummy_submodule)
+    # Is available here, but was not defined here
+    assert not is_native(dummy_submodule.getmembers, dummy_submodule)
+    # Argument 2 must be a module
+    with pytest.raises(TypeError):
+        is_native(dummy_submodule.some_function, "profile")
+    # Primitives have no source module information
+    with pytest.raises(TypeError):
+        is_native('text', dummy_submodule)
+
+
+def test_get_native_functions():
+    dummy_natives = get_native_functions(dummy_submodule)
+    # We know about these from the way the test is set up
+    assert dummy_submodule.some_function in dummy_natives
+    assert dummy_submodule.SomeClass.method in dummy_natives
+    assert dummy_submodule.SomeClass.cls_method in dummy_natives
+    assert dummy_submodule.SomeClass.stat_method in dummy_natives
+    # This shouldn't be there at all
+    assert test_get_native_functions not in dummy_natives
+    # This is imported in profile but not native
+    assert dummy_submodule.getmembers not in dummy_natives
+
+
+def test_get_submodules():
+    submodules = get_submodules('pcdsutils.tests')
+    assert dummy_submodule in submodules
+
+
+def test_basic_profiler():
+    # Run through and make sure our functions are included
+    with profiler_context(['pcdsutils']) as profiler:
+        dummy_submodule.some_function()
+        some_obj = dummy_submodule.SomeClass()
+        some_obj.method()
+        some_obj.cls_method()
+        some_obj.stat_method()
+
+    timings = profiler.get_stats().timings
+    functions_profiled = [
+        (os.path.basename(file), func)
+        for (file, lineno, func), stats in timings.items() if stats
+    ]
+    logger.debug(functions_profiled)
+    assert ('dummy_submodule.py', '__init__') in functions_profiled
+    assert ('dummy_submodule.py', 'some_function') in functions_profiled
+    assert ('dummy_submodule.py', 'method') in functions_profiled
+    assert ('dummy_submodule.py', 'cls_method') in functions_profiled
+    assert ('dummy_submodule.py', 'stat_method') in functions_profiled