Scope matching for LineProfiler.add_*()

line_profiler/autoprofile/line_profiler_utils.py[i]
    add_imported_function_or_module()
        Added the `match_scope` argument for limiting the scope of
        descension into classes in namespaces (classes and modules)

line_profiler/line_profiler.py[i]
    is_c_level_callable()
        New check for non-profilable C-level callables
    LineProfiler
        add_callable(), wrap_callable()
            Now no-ops on C-level callables
        add_class(), add_module()
            - Added the `match_scope` argument for limiting the scope of
              descension into classes in namespaces (classes and
              modules)
            - Added handling for when the `setattr()` on the namespace
              fails
        __call__()
            Added missing method in stub file

Author: TTsangSC

line_profiler/autoprofile/line_profiler_utils.py

@@ -1,7 +1,8 @@
 import inspect
 
 
-def add_imported_function_or_module(self, item, *, wrap=False):
+def add_imported_function_or_module(self, item, *,
+                                    match_scope='siblings', wrap=False):
     """
     Method to add an object to `LineProfiler` to be profiled.
 
@@ -12,6 +13,21 @@ def add_imported_function_or_module(self, item, *, wrap=False):
     Args:
         item (Union[Callable, Type, ModuleType]):
             Object to be profiled.
+        match_scope (Literal['exact', 'siblings', 'descendants',
+                             'none']):
+            Whether (and how) to match the scope of member classes to
+            `item` (if a class or module) and decide on whether to add
+            them:
+            - 'exact': only add classes defined locally in the body of
+              `item`
+            - 'descendants': only add locally-defined classes and
+              classes defined in submodules or locally-defined class
+              bodies, and so on.
+            - 'siblings': only add classes fulfilling 'descendants',
+              or defined in the same module as `item` (if a class) or in
+              sibling modules and subpackages to `item` (if a module)
+            - 'none': don't check scopes and add all classes in the
+              namespace
         wrap (bool):
             Whether to replace the wrapped members with wrappers which
             automatically enable/disable the profiler when called.
@@ -23,9 +39,9 @@ def add_imported_function_or_module(self, item, *, wrap=False):
         `LineProfiler.add_callable()`, `.add_module()`, `.add_class()`
     """
     if inspect.isclass(item):
-        count = self.add_class(item, wrap=wrap)
+        count = self.add_class(item, match_scope=match_scope, wrap=wrap)
     elif inspect.ismodule(item):
-        count = self.add_module(item, wrap=wrap)
+        count = self.add_module(item, match_scope=match_scope, wrap=wrap)
     else:
         try:
             count = self.add_callable(item)

line_profiler/autoprofile/line_profiler_utils.pyi

@@ -1,8 +1,21 @@
-from typing import Callable, Literal
 from types import ModuleType
+from typing import overload, Any, Literal, TYPE_CHECKING
 
+if TYPE_CHECKING:  # Stub-only annotations
+    from ..line_profiler import CLevelCallable, CallableLike, MatchScopeOption
 
+
+@overload
+def add_imported_function_or_module(
+        self, item: CLevelCallable | Any,
+        match_scope: MatchScopeOption = 'siblings',
+        wrap: bool = False) -> Literal[0]:
+    ...
+
+
+@overload
 def add_imported_function_or_module(
-        self, item: Callable | type | ModuleType, *,
+        self, item: CallableLike | type | ModuleType,
+        match_scope: MatchScopeOption = 'siblings',
         wrap: bool = False) -> Literal[0, 1]:
     ...

line_profiler/line_profiler.py

@@ -10,6 +10,8 @@
 import tempfile
 import os
 import sys
+import types
+import warnings
 from argparse import ArgumentError, ArgumentParser
 
 try:
@@ -28,9 +30,28 @@
 # NOTE: This needs to be in sync with ../kernprof.py and __init__.py
 __version__ = '4.3.0'
 
+# These objects are callables, but are defined in C so we can't handle
+# them anyway
+c_level_callable_types = (types.BuiltinFunctionType,
+                          types.BuiltinMethodType,
+                          types.ClassMethodDescriptorType,
+                          types.MethodDescriptorType,
+                          types.MethodWrapperType,
+                          types.WrapperDescriptorType)
+
 is_function = inspect.isfunction
 
 
+def is_c_level_callable(func):
+    """
+    Returns:
+        func_is_c_level (bool):
+            Whether a callable is defined at the C level (and is thus
+            non-profilable).
+    """
+    return isinstance(func, c_level_callable_types)
+
+
 def load_ipython_extension(ip):
     """ API for IPython to recognize this module as an IPython extension.
     """
@@ -63,6 +84,8 @@ def _get_underlying_functions(func):
                         f'cannot get functions from {type(func)} objects')
     if is_function(func):
         return [func]
+    if is_c_level_callable(func):
+        return []
     return [type(func).__call__]
 
 
@@ -90,12 +113,20 @@ def __call__(self, func):
         start the profiler on function entry and stop it on function
         exit.
         """
-        # Note: if `func` is a `types.FunctionType` which is already
-        # decorated by the profiler, the same object is returned;
+        # The same object is returned when:
+        # - `func` is a `types.FunctionType` which is already
+        #   decorated by the profiler, or
+        # - `func` is any of the C-level callables that can't be
+        #   profiled
         # otherwise, wrapper objects are always returned.
         self.add_callable(func)
         return self.wrap_callable(func)
 
+    def wrap_callable(self, func):
+        if is_c_level_callable(func):  # Non-profilable
+            return func
+        return super().wrap_callable(func)
+
     def add_callable(self, func):
         """
         Register a function, method, property, partial object, etc. with
@@ -132,50 +163,153 @@ def print_stats(self, stream=None, output_unit=None, stripzeros=False,
                   stream=stream, stripzeros=stripzeros,
                   details=details, summarize=summarize, sort=sort, rich=rich)
 
-    def _add_namespace(self, namespace, *, wrap=False):
-        """
-        Add the members (callables (wrappers), methods, classes, ...) in
-        a namespace and profile them.
-
-        Args:
-            namespace (Union[ModuleType, type]):
-                Module or class to be profiled.
-            wrap (bool):
-                Whether to replace the wrapped members with wrappers
-                which automatically enable/disable the profiler when
-                called.
-
-        Returns:
-            n (int):
-                Number of members added to the profiler.
-        """
-        return self._add_namespace_inner(set(), namespace, wrap=wrap)
-
-    def _add_namespace_inner(self, duplicate_tracker, namespace, *,
-                             wrap=False):
+    def _add_namespace(self, duplicate_tracker, namespace, *,
+                       filter_scope=None, wrap=False):
         count = 0
-        add_cls = self._add_namespace_inner
+        add_cls = self._add_namespace
         add_func = self.add_callable
+        wrap_failures = {}
+        if filter_scope is None:
+            def filter_scope(*_):
+                return True
+
         for attr, value in vars(namespace).items():
             if id(value) in duplicate_tracker:
                 continue
             duplicate_tracker.add(id(value))
             if isinstance(value, type):
-                if add_cls(duplicate_tracker, value, wrap=wrap):
-                    count += 1
+                if filter_scope(namespace, value):
+                    if add_cls(duplicate_tracker, value, wrap=wrap):
+                        count += 1
                 continue
             try:
-                func_needs_adding = add_func(value)
+                if not add_func(value):
+                    continue
             except TypeError:  # Not a callable (wrapper)
                 continue
-            if not func_needs_adding:
-                continue
             if wrap:
-                setattr(namespace, attr, self.wrap_callable(value))
+                wrapper = self.wrap_callable(value)
+                if wrapper is not value:
+                    try:
+                        setattr(namespace, attr, wrapper)
+                    except (TypeError, AttributeError):
+                        # Corner case in case if a class/module don't
+                        # allow setting attributes (could e.g. happen
+                        # with some builtin/extension classes, but their
+                        # method should be in C anyway, so
+                        # `.add_callable()` should've returned 0 and we
+                        # shouldn't be here)
+                        wrap_failures[attr] = value
             count += 1
+        if wrap_failures:
+            msg = (f'cannot wrap {len(wrap_failures)} attribute(s) of '
+                   f'{namespace!r} (`{{attr: value}}`): {wrap_failures!r}')
+            warnings.warn(msg, stacklevel=2)
         return count
 
-    add_class = add_module = _add_namespace
+    def add_class(self, cls, *, match_scope='siblings', wrap=False):
+        """
+        Add the members (callables (wrappers), methods, classes, ...) in
+        a class' local namespace and profile them.
+
+        Args:
+            cls (type):
+                Class to be profiled.
+            match_scope (Literal['exact', 'siblings', 'descendants',
+                                 'none']):
+                Whether (and how) to match the scope of member classes
+                and decide on whether to add them:
+                - 'exact': only add classes defined locally in this
+                  namespace, i.e. in the body of `cls`, as "inner
+                  classes"
+                - 'descendants': only add "inner classes", their "inner
+                  classes", and so on.
+                - 'siblings': only add classes fulfilling 'descendants',
+                  or defined in the same module as `cls`
+                - 'none': don't check scopes and add all classes in the
+                  namespace
+            wrap (bool):
+                Whether to replace the wrapped members with wrappers
+                which automatically enable/disable the profiler when
+                called.
+
+        Returns:
+            n (int):
+                Number of members added to the profiler.
+        """
+        def class_is_child(cls, other):
+            if not modules_are_equal(cls, other):
+                return False
+            return other.__qualname__ == f'{cls.__qualname__}.{other.__name__}'
+
+        def modules_are_equal(cls, other):  # = sibling check
+            return cls.__module__ == other.__module__
+
+        def class_is_descendant(cls, other):
+            if not modules_are_equal(cls, other):
+                return False
+            return other.__qualname__.startswith(cls.__qualname__ + '.')
+
+        filter_scope = {'exact': class_is_child,
+                        'descendants': class_is_descendant,
+                        'siblings': modules_are_equal,
+                        'none': None}[match_scope]
+        return self._add_namespace(set(), cls,
+                                   filter_scope=filter_scope, wrap=wrap)
+
+    def add_module(self, mod, *, match_scope='siblings', wrap=False):
+        """
+        Add the members (callables (wrappers), methods, classes, ...) in
+        a module's local namespace and profile them.
+
+        Args:
+            mod (ModuleType):
+                Module to be profiled.
+            match_scope (Literal['exact', 'siblings', 'descendants',
+                                 'none']):
+                Whether (and how) to match the scope of member classes
+                and decide on whether to add them:
+                - 'exact': only add classes defined locally in this
+                  namespace, i.e. in the body of `mod`
+                - 'descendants': only add locally-defined classes,
+                  classes locally defined in their bodies, and so on
+                - 'siblings': only add classes fulfilling 'descendants',
+                  or defined in sibling modules/subpackages to `mod` (if
+                  `mod` is part of a package)
+                - 'none': don't check scopes and add all classes in the
+                  namespace
+            wrap (bool):
+                Whether to replace the wrapped members with wrappers
+                which automatically enable/disable the profiler when
+                called.
+
+        Returns:
+            n (int):
+                Number of members added to the profiler.
+        """
+        def match_prefix(s: str, prefix: str, sep: str = '.') -> bool:
+            return s == prefix or s.startswith(prefix + sep)
+
+        def class_is_child(mod, other):
+            return other.__module__ == mod.__name__
+
+        def class_is_descendant(mod, other):
+            return match_prefix(other.__module__, mod.__name__)
+
+        def class_is_cousin(mod, other):
+            if class_is_descendant(mod, other):
+                return True
+            return match_prefix(other.__module__, parent)
+
+        parent, _, basename = mod.__name__.rpartition('.')
+        filter_scope = {'exact': class_is_child,
+                        'descendants': class_is_descendant,
+                        'siblings': (class_is_cousin  # Only if a pkg
+                                     if basename else
+                                     class_is_descendant),
+                        'none': None}[match_scope]
+        return self._add_namespace(set(), mod,
+                                   filter_scope=filter_scope, wrap=wrap)
 
 
 # This could be in the ipython_extension submodule,

line_profiler/line_profiler.pyi

@@ -1,16 +1,51 @@
-from types import ModuleType
-from typing import List, Literal, Tuple
+from functools import cached_property, partial, partialmethod
+from inspect import isfunction as is_function
+from types import (FunctionType, MethodType, ModuleType,
+                   BuiltinFunctionType, BuiltinMethodType,
+                   ClassMethodDescriptorType, MethodDescriptorType,
+                   MethodWrapperType, WrapperDescriptorType)
+from typing import (overload,
+                    Any, Callable, List, Literal, Tuple, TypeGuard, TypeVar)
 import io
 from ._line_profiler import LineProfiler as CLineProfiler
 from .profiler_mixin import ByCountProfilerMixin
 from _typeshed import Incomplete
 
 
+CLevelCallable = TypeVar('CLevelCallable',
+                         BuiltinFunctionType, BuiltinMethodType,
+                         ClassMethodDescriptorType, MethodDescriptorType,
+                         MethodWrapperType, WrapperDescriptorType)
+CallableLike = TypeVar('CallableLike',
+                       FunctionType, partial, property, cached_property,
+                       MethodType, staticmethod, classmethod, partialmethod)
+MatchScopeOption = Literal['exact', 'descendants', 'siblings', 'none']
+
+
+def is_c_level_callable(func: Any) -> TypeGuard[CLevelCallable]:
+    ...
+
+
 def load_ipython_extension(ip) -> None:
     ...
 
 
 class LineProfiler(CLineProfiler, ByCountProfilerMixin):
+    @overload
+    def __call__(self,  # type: ignore[overload-overlap]
+                 func: CLevelCallable) -> CLevelCallable:
+        ...
+
+    @overload
+    def __call__(self,  # type: ignore[overload-overlap]
+                 func: CallableLike) -> CallableLike:
+        ...
+
+    # Fallback: just wrap the `.__call__()` of a generic callable
+
+    @overload
+    def __call__(self, func: Callable) -> FunctionType:
+        ...
 
     def add_callable(self, func) -> Literal[0, 1]:
         ...
@@ -28,10 +63,14 @@ class LineProfiler(CLineProfiler, ByCountProfilerMixin):
                     rich: bool = ...) -> None:
         ...
 
-    def add_module(self, mod: ModuleType, *, wrap: bool = False) -> int:
+    def add_module(self, mod: ModuleType, *,
+                   match_scope: MatchScopeOption = 'siblings',
+                   wrap: bool = False) -> int:
         ...
 
-    def add_class(self, cls: type, *, wrap: bool = False) -> int:
+    def add_class(self, cls: type, *,
+                  match_scope: MatchScopeOption = 'siblings',
+                  wrap: bool = False) -> int:
         ...