python cprofile decorator

Question

If you're anything like me and sometimes write pretty slow code... you may end up going back to API's or scripts and wanting to profile an endpoint/method to see what is taking so long - so I can focus on the areas to refactor improvements into.

There is a ton of profiling packages and content out there for python. But every single time, I end up reading the same terrible tutorials on how to profile with CProfile. I like it a lot, but I hate importing Pstats and other packages, indenting my code and surrounding everything with the correct calls to produce a simple profile output.

Hence, I made a decorator and uploaded it to pip: https://pypi.org/project/profiling-decorator/

I'm looking for feedback as to whether this is useful at all, or if there is just something much easier out there I could be using instead. I uploaded it to pip more as practice for me to understand how the workflow processes should be, as I am new to all of it.

Source code:

"""profiling decorator definition"""

# pylint: disable = missing-function-docstring
import asyncio
import cProfile
from functools import wraps
import io
import logging
import pstats
from pstats import SortKey
from typing import Any, Callable, List, Optional, Union, Tuple


VALID_SORTS = set(
    val
    for name, attribute in SortKey.__dict__.items()
    if not name.startswith("__") and isinstance(attribute, tuple)
    for val in attribute
)


def profile(
    n_rows: int = 50,
    sort_by: Union[str, List[str], Tuple] = "cumulative",
    output: str = "stdout",
    filename: Optional[str] = None,
) -> Callable:
    """
    Decorator to profile a function.

    Args:
        n_rows: Number of rows to reduce profile output to.
        sort_by: What statistic(s) to sort by in the results.
        output: How to output the results.
        filename: The filename to output to if output='file'

    Returns:
        func return

    Raises:
        ValueError: Invalid value for parameter.
        IOError: Error writing to file.
    """
    valid_outputs = {"stdout", "file", "log"}

    def decorator(func: Callable) -> Callable:

        @wraps(func)
        async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
            with cProfile.Profile() as pr:
                # Directly await the async function here
                result = await func(*args, **kwargs)
                process_profiling_results(pr)

            return result

        @wraps(func)
        def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
            with cProfile.Profile() as pr:
                result = func(*args, **kwargs)
                process_profiling_results(pr)

            return result

        if asyncio.iscoroutinefunction(func):
            return async_wrapper
        return sync_wrapper

    def process_profiling_results(pr: cProfile.Profile) -> None:
        s = io.StringIO()
        # Fallback to 'cumulative' if sort_by is invalid
        if isinstance(sort_by, (list, tuple)):
            sort_by_valid = [s for s in sort_by if s in VALID_SORTS]
        else:
            sort_by_valid = [sort_by] if sort_by in VALID_SORTS else ["cumulative"]

        ps = pstats.Stats(pr, stream=s).sort_stats(*sort_by_valid)
        ps.print_stats(n_rows)

        # Handle output
        if output not in valid_outputs:
            raise ValueError(
                f"Invalid output option '{output}'. Valid options are {valid_outputs}."
            )

        if output == "stdout":
            print(s.getvalue())
        elif output == "file":
            if not filename:
                raise ValueError("Filename must be provided when output is 'file'.")
            try:
                with open(filename, "w+", encoding="utf-8") as f:
                    f.write(s.getvalue())
            except IOError as e:
                raise IOError(f"Error writing to file {filename}: {e}") from e
        elif output == "log":
            logger = logging.getLogger(__name__)
            logger.info(s.getvalue())
        else:
            raise ValueError("'output' must be one of 'stdout', 'file' or 'log'.")

    # Enable the decorator to be used without parentheses if no arguments are provided
    if callable(n_rows):
        temp_func = n_rows
        n_rows = 50
        return decorator(temp_func)

    return decorator

Answer 1

Overall, I feel this package makes a helpful contribution, for which I thank you. It eases the process of producing meaningful profile dumps.

ancient interpreters

Supporting downrev cPythons, such as 3.8, is a laudable goal. We see it mentioned in the project file:

requires-python = ">=3.8"

Consider adding this within the src/ or tests/ directories, to make it more explicit so we may get a runtime failure:

import sys

assert sys.version_info >= (3, 8)

This helps us understand the outdated Union[x, y] annotations, where modern code would prefer x | y.

Come November of this year, once 3.8 is past EOL, we would expect that to bump up to 3.9. At that point fatal asserts triggering under 3.8 would offer helpful diagnostic advice about the toolchain support situation, encouraging updates to better supported code.

I do really like those helpful annotations, even though parts of the Typing module will be phased out before long.

CI / CD

For the five current interpreters you're targetting, 3.8 .. 3.12, it would be useful to set up a GitHub workflow or similar cron job that tests your released package against each interpreter. That would help instill confidence that a given maintenance edit didn't introduce lossage in an old or new variant, e.g. due to cprofile module changes.

test suite

I am delighted to see some smoke tests in the repo.

They would benefit from some code that makes a cursory inspection of the profiling output and verifies that it changes when the requested sort order changes. For example, a test routine might repeatedly invoke foo() and bar(), with one accounting for more calls and the other accounting for more elapsed time.

lint

Thank you for having run black on this. Please run $ isort .

# pylint: disable = missing-function-docstring

Consider eliding this line, as it is redundant with the project config:

[tool.pylint]
...
disable = [
    "C0114", # (missing-module-docstring)
    "C0115", # (missing-class-docstring)
    "C0116", # (missing-function-docstring)
]

Those comments BTW are perfect!

perils of supporting multiple versions

VALID_SORTS = set(
    val
    for name, attribute in SortKey.__dict__.items()
    if not name.startswith("__") and isinstance(attribute, tuple)
    for val in attribute
)

I don't know what that means; it would benefit from an explanatory """docstring""" which maybe gives an example or two. I only have a 3.12.2 interpreter handy, and there it reports empty-set. Your "pick out the tuples" logic is too fragile.

I imagine you're looking for something like this:

>>> SortKey.__dict__['_member_names_']
['CALLS', 'CUMULATIVE', 'FILENAME', 'LINE', 'NAME', 'NFL', 'PCALLS', 'STDNAME', 'TIME']
>>>
>>> sorted(set(map(str, SortKey)))
['calls', 'cumulative', 'filename', 'line', 'name', 'nfl', 'pcalls', 'stdname', 'time']

Not sure which. The _member_names_ approach is not a public interface.

The Fallback to 'cumulative' behavior regrettably masks this for default runs, revealing the trouble only if caller e.g. specifies 'calls'.

Moral of the story is: Come in through the (supported) front door. If you use the back door, be sure you know what you're doing, in all contexts.

annotate generics

I don't understand this line:

def profile(
    ...
    sort_by: Union[str, List[str], Tuple] = ...

Prefer:

    sort_by: Union[str, List[str], Tuple[str, ...]] = ...

But if we're gonna go there, wouldn't it make more sense to just talk about a Sequence[str] of sort keys?

pathlib

    filename: Optional[str] = None,

Prefer a Path here, or perhaps Union[str, Path].

Also the Callable annotation omits [..., Any], but maybe that's a 3.8 thing. Not sure if 3.9 started to support the in + out types.

stringly typed

    output: str = "stdout",

That's odd. Consider accepting sys.stdout and then doing the Right Thing with it. Or at least use an enum, to replace valid_outputs.

Using two parameters to represent one thing, a destination, seems odd. Needing to specify output="file" seems unnatural.

Inability to send output to sys.stderr seems very odd; we can't invoke python script.py > results.txt 2> profile.txt

It's weird that I can't pass in my own pre-configured logger.

Kudos for importing wraps! Nice use of foo_wrappers.

Consider doing Extract Helper for the # Handle output section.

no parens

Kudos for flexibly supporting both @profile and @profile(). However, this should be DRYed up:

    # Enable the decorator to be used without parentheses if no arguments are provided
    if callable(n_rows):
        temp_func = n_rows
        n_rows = 50
        return decorator(temp_func)

At the very least define DEFAULT_ROWS = 50, as we also use that magic number up in the signature.

I don't know the Right Way to offer such flexible support, but I feel that this is not it. I confess that a cursory review of the LRU cache source code seems pretty similar to your code, and perhaps that's not coincidental. Please suggest other examples of such usage that we might critique together. It seems like a moderately common usage for which there should be a nicer solution.

---

This submission substantially accomplishes its design goals.

I would be willing to delegate or accept maintenance tasks on it.

Answer 2

Type Hinting Literal Values

You only accept stdout, file, or log for your output values, but you wouldn't know that from the type hints or docstrings, only from an error being raised. Let's address both:

from typing import Any, Callable, List, Literal, Optional, Union, Tuple

def profile(
    n_rows: int = 50,
    sort_by: Union[str, List[str], Tuple] = "cumulative",
    output: Literal["stdout", "file", "log"] = "stdout",
    filename: Optional[str] = None,
) -> Callable:
    """
    Decorator to profile a function.

    Args:
        n_rows: Number of rows to reduce profile output to.
        sort_by: What statistic(s) to sort by in the results.
        output: How to output the results. Should be either 
            "stdout", "file", or "log"
        filename: The filename to output to if output='file'

    Returns:
        func return

    Raises:
        ValueError: Invalid value for stdout.
        IOError: Error writing to file.
    """
    # rest of function

Unreachable else

You have a check for if not output in valid_outputs, which means that in the if block that follows, the else cannot possibly be reached:

# Handle output
if output not in valid_outputs:
    raise ValueError(
        f"Invalid output option '{output}'. Valid options are {valid_outputs}."
    )

if output == "stdout":
    print(s.getvalue())
elif output == "file":
    ...
elif output == "log":
    ...
# This isn't reachable, you would have raised the error already
else:
    raise ValueError("'output' must be one of 'stdout', 'file' or 'log'.")

In this case, the else should replace the elif output == "log" and the actual check for output value should be at the top of the decorator. Raise the error early, rather than waiting to profile the entire (possibly long-running) function to find out the decorator was configured incorrectly:

def profile(...):
    if output not in ('stdout', 'file', 'log'):
        raise ValueError(...)
    # while we're here, let's knock this one out, too
    elif output == "file" and filename is None:
        raise ValueError(...)

if callable

I think this is breaking the contract defined by your function signature. Let's add an argument for a Callable that defaults to None:

from functools import partial 

def profile(
    func: Optional[Callable]=None,
    n_rows: int = 50,
    sort_by: Union[str, List[str], Tuple] = "cumulative",
    output: Literal["stdout", "file", "log"] = "stdout",
    filename: Optional[str] = None,
) -> Callable:
    if func is None:
        # return a function that is ready to accept the callable as
        # an argument
        return partial(
            profile, 
            n_rows=n_rows, 
            sort_by=sort_by, 
            output=output, 
            filename=filename
        )
    
    # rest of function

Answer 3

Both answers posted so far touch type hints, but do not cover the issues in full. Let's start with them.

Type hints

You support decorating with both @profile and @profile(...). However, your function is hinted to return just Callable. Type parameters in python are always defaulted to Any, so this is equivalent to Callable[..., Any] - a callable that takes any args and kwargs and returns the most permissive type - Any. It's bad for users:

@profile
def my_func(x: int) -> None: 
    return None

my_func(foo="bar")  # No typechecker error!

Let's fix that. First I will show py3.8 native compatible method - it does not give full type safety, but makes your interface correct for an external observer. Note Tuple fixes - see @J_H answer.

from typing import Any, Callable, List, Optional, Union, Tuple, TypeVar, overload

_C = TypeVar("_C", bound=Callable[..., Any])

@overload
def profile(function: _C, /) -> _C: ...
@overload
def profile(
    n_rows: int = 50,
    sort_by: Union[str, List[str], Tuple[str, ...]] = "cumulative",
    output: str = "stdout",
    filename: Optional[str] = None,
) -> Callable[[_C], _C]: ...


def profile(
    n_rows: int | Callable[..., Any] = 50,
    sort_by: Union[str, List[str], Tuple[str, ...]] = "cumulative",
    output: str = "stdout",
    filename: Optional[str] = None,
) -> Callable[..., Any]:
    ...  # Your function body goes here

What's going on? We declare a function that can either be called with one callable argument or with multiple non-callable arguments. overloaded definitions are (almost) ignored at runtime, they have no effect on behaviour - your function's signature is still the one not decorated with overload. On the other hand, typecheckers will now know that @profile returns function unchanged, and @profile(...) also does (because Callable[[_C], _C] is what is applied to the decorated function - a callable that returns its only argument unchanged). Great. Now the first snippet in this answer will cause a typechecker error.

You can see this solution working in playground to understand these typing details better.

But we have too much freedom! Our implementation is still too permissive inside: type checker won't flag many errors if they are there. Your version limitation comes now into play. If you need to be compatible with older versions (which is a good thing for libraries), you can add a lightweight dependency - typing_extensions. This package is maintained by mypy and python core team and backports new typing features to make them available on older pythons. Once you increase minimal supported versions, some imports should be moved to regular typing (but you can leave them as is, typing_extensions is guaranteed to be backwards-compatible). pyupgrade and ruff's UP* rules will catch that too.

I will show the "final" (to the best of my knowledge) solution with most typechecking you can get below and try to explain it a bit.

# non-typing import unchanged
from typing import TYPE_CHECKING, Any, Awaitable, Callable, List, Optional, Union, Tuple, TypeVar, cast, overload
from typing_extensions import ParamSpec

_C = TypeVar("_C", bound=Callable[..., Any])
_R = TypeVar("_R")
_P = ParamSpec("_P")

@overload
def profile(function: _C, /) -> _C: ...
@overload
def profile(
    n_rows: int = 50,
    sort_by: Union[str, List[str], Tuple[str, ...]] = "cumulative",
    output: str = "stdout",
    filename: Optional[str] = None,
) -> Callable[[_C], _C]: ...

def profile(
    n_rows: int | Callable[_P, _R] = 50,
    sort_by: Union[str, List[str], Tuple[str, ...]] = "cumulative",
    output: str = "stdout",
    filename: Optional[str] = None,
) -> Callable[_P, _R] | Callable[[Callable[_P, _R]], Callable[_P, _R]]:
    # Your docstring here
    valid_outputs = {"stdout", "file", "log"}

    def decorator(func: Callable[_P, _R]) -> Callable[_P, _R]:
        if asyncio.iscoroutinefunction(func):
            # Now it's an async callable. Reuse `_R`, but inside `Awaitable`
            # now we can await it's return type
            afunc = cast("Callable[_P, Awaitable[_R]]", func)
            
            @wraps(afunc)
            async def async_wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R:
                with cProfile.Profile() as pr:
                    # Directly await the async function here
                    result = await afunc(*args, **kwargs)
                    process_profiling_results(pr)
    
                return result
            
            # async functions are interpreted as returning `Awaitable[DeclaredReturnType]`,
            # we cannot get rid of this onion layer
            return async_wrapper  # type: ignore[return-value]

        @wraps(func)
        def sync_wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R:
            with cProfile.Profile() as pr:
                result = func(*args, **kwargs)
                process_profiling_results(pr)

            return result

        return sync_wrapper

    def process_profiling_results(pr: cProfile.Profile) -> None:
        s = io.StringIO()
        # Fallback to 'cumulative' if sort_by is invalid
        if isinstance(sort_by, (list, tuple)):
            sort_by_valid = [s for s in sort_by if s in VALID_SORTS]
        else:
            sort_by_valid = [sort_by] if sort_by in VALID_SORTS else ["cumulative"]

        ps = pstats.Stats(pr, stream=s).sort_stats(*sort_by_valid)
        assert isinstance(n_rows, int)  # Due to our trickery with the first argument
        ps.print_stats(n_rows)

        # Handle output
        if output not in valid_outputs:
            raise ValueError(
                f"Invalid output option '{output}'. Valid options are {valid_outputs}."
            )

        if output == "stdout":
            print(s.getvalue())
        elif output == "file":
            if not filename:
                raise ValueError("Filename must be provided when output is 'file'.")
            try:
                with open(filename, "w+", encoding="utf-8") as f:
                    f.write(s.getvalue())
            except IOError as e:
                raise IOError(f"Error writing to file {filename}: {e}") from e
        elif output == "log":
            logger = logging.getLogger(__name__)
            logger.info(s.getvalue())
        else:
            raise ValueError("'output' must be one of 'stdout', 'file' or 'log'.")

    # Enable the decorator to be used without parentheses if no arguments are provided
    if callable(n_rows):
        temp_func = n_rows
        n_rows = 50
        return decorator(temp_func)

    return decorator

Now (again, playground) this would've caught, for example, forgetting to pass *args down to the callable. Is it worth the tradeoff? Not sure, typing sync- and async-enabled decorators is always a terrible mess, and I usually prefer just having two different decorators (then both can be easily type annotated). However, the main goal is still achieved: external observers have their function's signatures unchanged.

Now, other minor stuff (I'm ignoring everything already mentioned).

list | tuple versus any sequence or iterable

Typing aside, accepting only a list or tuple is not really user-friendly. There are other sequences and iterables. The only thing you do with this argument is unpacking/iterating, so Iterable would be the most broad type possible and require only minor modifications from you. On second thought I think Iterable is a wrong choice, because Iterable does not guarantee fixed/known iteration order, and order matters in this case. Sequence is a random-access iterable with integer indices, so it's pretty much ordered. Let's replace your list-or-tuple with a sequence:

...
from typing import Sequence

# Correct the overload too, omitted here for brewity
def profile(
    n_rows: int | Callable[_P, _R] = 50,
    sort_by: Union[str, Sequence[str]] = "cumulative",
    output: str = "stdout",
    filename: Optional[str] = None,
) -> Callable[_P, _R] | Callable[[Callable[_P, _R]], Callable[_P, _R]]:
    ...
    def process_profiling_results(pr: cProfile.Profile) -> None:
        if isinstance(sort_by, str):
            sort_by_valid = [sort_by] if sort_by in VALID_SORTS else ["cumulative"]
        else:
            sort_by_valid = [s for s in sort_by if s in VALID_SORTS]

Note that str is a subtype of Sequence[str] so we can't isinstance(sort_by, Sequence) - just swap the branches. Union[str, Sequence[str]] is strictly the same as just Sequence[str], but I'd rather keep the union for human readers.

Do not fail silently

Your sort_by fallback may be convenient for somebody, but a debugging disaster for everybody else. It's better to raise an exception if you encounter an invalid sort key - this way a typo will be easy to catch.

Set comprehension

VALID_SORTS = set(
    val
    for name, attribute in SortKey.__dict__.items()
    if not name.startswith("__") and isinstance(attribute, tuple)
    for val in attribute
)

is... interesting. Why is there a generator wrapped with set() function? Python has set comprehension literals.

VALID_SORTS = {
    val
    for name, attribute in SortKey.__dict__.items()
    if not name.startswith("__") and isinstance(attribute, tuple)
    for val in attribute
}

Enum members retrieval (already covered, sorry)

pstats.SortKey is a enum.Enum. Please don't do what you're doing with it. You can iterate over enum directly to get its members. Furthermore, your current solution results in empty VALID_SORTS for me on py3.11. And on py3.8. Looks like this doesn't work on any python, actually. You probably just need list(map(str, SortKey)) or [str(key) for key in SortKey], whichever reads better for you - this will collect all lowercase key versions.

But now... you don't need VALID_SORTS at all. To validate, SortKey(key) will raise a ValueError if key is invalid.

Now you're at the point where you have replaced an enum with its stringified members. Why? Same "stringly typed" issue. I'd recommend to ditch that and accept SortKey | Sequence[SortKey]. You can even reexport SortKey in your package to prevent using several imports - that's perfectly fine for an utility package, so that it's just from your_package import SortKey, profile for your users.