Skip to content

Add ok/error logic to the decorator #79

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 19 commits into from
Aug 28, 2023
Merged
Show file tree
Hide file tree
Changes from 16 commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 5 additions & 3 deletions .github/workflows/main.yml
Original file line number Diff line number Diff line change
Expand Up @@ -24,10 +24,12 @@ jobs:
python-version: ${{ matrix.python-version }}
cache: poetry
- name: Install dependencies
run: poetry install --no-interaction --no-root --with dev
run: poetry install --no-interaction --no-root --with dev --with examples
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

--with supports comma separated syntax too 🙃

Suggested change
run: poetry install --no-interaction --no-root --with dev --with examples
run: poetry install --no-interaction --no-root --with dev,examples

- name: Check code formatting
run: poetry run black .
- name: Lint code
run: poetry run pyright
- name: Lint lib code
run: poetry run mypy src --enable-incomplete-feature=Unpack
- name: Lint lib examples
run: poetry run mypy examples --enable-incomplete-feature=Unpack
- name: Run tests
run: poetry run pytest -n auto
6 changes: 3 additions & 3 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -162,7 +162,7 @@ poetry install --with examples
Code in this repository is:

- formatted using [black](https://black.readthedocs.io/en/stable/).
- contains type definitions (which are linted by [pyright](https://microsoft.github.io/pyright/))
- contains type definitions (which are linted by [mypy](https://www.mypy-lang.org/))
- tested using [pytest](https://docs.pytest.org/)

In order to run these tools locally you have to install them, you can install them using poetry:
Expand All @@ -176,8 +176,8 @@ After that you can run the tools individually
```sh
# Formatting using black
poetry run black .
# Lint using pyright
poetry run pyright
# Lint using mypy
poetry run mypy .
# Run the tests using pytest
poetry run pytest
# Run a single test, and clear the cache
Expand Down
6 changes: 6 additions & 0 deletions examples/django_example/mypy.ini
Original file line number Diff line number Diff line change
@@ -0,0 +1,6 @@
[mypy]
plugins =
mypy_django_plugin.main

[mypy.plugins.django-stubs]
django_settings_module = "django_example.settings"
468 changes: 329 additions & 139 deletions poetry.lock

Large diffs are not rendered by default.

26 changes: 24 additions & 2 deletions pyproject.toml
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

since we are exposing types correctly now, i think it would be fair to also add the "Typing :: Typed" classifier 🧐

Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,6 @@ classifiers = [
packages = [{ include = "autometrics", from = "src" }]

[tool.poetry.dependencies]
opentelemetry-api = "^1.17.0"
opentelemetry-exporter-prometheus = "^1.12.0rc1"
opentelemetry-sdk = "^1.17.0"
prometheus-client = "^0.16.0 || ^0.17.0"
Expand All @@ -35,12 +34,28 @@ typing-extensions = "^4.5.0"
[tool.poetry.group.dev]
optional = true

[tool.mypy]
namespace_packages = true
mypy_path = "src"
enable_incomplete_feature = "Unpack"

# This override is needed because with certain flavors of python and
# mypy you can get the following error:
# opentelemetry/attributes/__init__.py:14: error: invalid syntax
# Which at the time of writing is a line that states ignore types:
# `# type: ignore`
[[tool.mypy.overrides]]
module = [
"opentelemetry.attributes",
]
follow_imports = "skip"

[tool.poetry.group.dev.dependencies]
pyright = "^1.1.307"
pytest = "^7.3.0"
pytest-asyncio = "^0.21.0"
black = "^23.3.0"
pytest-xdist = "^3.3.1"
mypy = "^1.5.1"
twine = "4.0.2"

[tool.poetry.group.examples]
Expand Down Expand Up @@ -83,6 +98,13 @@ uvicorn = "0.21.1"
webencodings = "0.5.1"
zipp = "3.15.0"
locust = "^2.15.1"
django-stubs = "4.2.3"



[tool.poetry.group.development.dependencies]
types-requests = "^2.31.0.2"
django-stubs = "^4.2.3"

[build-system]
requires = ["poetry-core"]
Expand Down
157 changes: 107 additions & 50 deletions src/autometrics/decorator.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@

from contextvars import ContextVar, Token
from functools import wraps
from typing import overload, TypeVar, Callable, Optional, Awaitable
from typing import overload, TypeVar, Callable, Optional, Awaitable, Union, Coroutine
from typing_extensions import ParamSpec

from .objectives import Objective
Expand All @@ -15,39 +15,60 @@
append_docs_to_docstring,
)


P = ParamSpec("P")
T = TypeVar("T")

Params = ParamSpec("Params")
R = TypeVar("R")
Y = TypeVar("Y")
S = TypeVar("S")

caller_module_var: ContextVar[str] = ContextVar("caller.module", default="")
caller_function_var: ContextVar[str] = ContextVar("caller.function", default="")


# Bare decorator usage
# Decorator with arguments (where decorated function returns an awaitable)
@overload
def autometrics(func: Callable[P, T]) -> Callable[P, T]:
def autometrics(
func: None = None,
*,
objective: Optional[Objective] = None,
track_concurrency: Optional[bool] = False,
record_error_if: Callable[[R], bool],
record_success_if: Optional[Callable[[Exception], bool]] = None,
) -> Union[
Callable[
[Callable[Params, Coroutine[Y, S, R]]], Callable[Params, Coroutine[Y, S, R]]
],
Callable[[Callable[Params, R]], Callable[Params, R]],
]:
...


# Decorator with arguments (where decorated function returns an awaitable)
@overload
def autometrics(func: Callable[P, Awaitable[T]]) -> Callable[P, Awaitable[T]]:
def autometrics(
func: None = None,
*,
objective: Optional[Objective] = None,
track_concurrency: Optional[bool] = False,
record_success_if: Optional[Callable[[Exception], bool]] = None,
) -> Callable[[Callable[Params, R]], Callable[Params, R]]:
...


# Decorator with arguments
# Using the func parameter
# i.e. using @autometrics()
@overload
def autometrics(
*, objective: Optional[Objective] = None, track_concurrency: Optional[bool] = False
) -> Callable:
func: Callable[Params, R],
) -> Callable[Params, R]:
...


def autometrics(
func: Optional[Callable] = None,
*,
objective: Optional[Objective] = None,
track_concurrency: Optional[bool] = False,
func=None,
objective=None,
track_concurrency=None,
record_error_if=None,
record_success_if=None,
):
"""Decorator for tracking function calls and duration. Supports synchronous and async functions."""

Expand Down Expand Up @@ -100,15 +121,15 @@ def track_result_error(
result=Result.ERROR,
)

def sync_decorator(func: Callable[P, T]) -> Callable[P, T]:
def sync_decorator(func: Callable[Params, R]) -> Callable[Params, R]:
"""Helper for decorating synchronous functions, to track calls and duration."""

module_name = get_module_name(func)
func_name = get_function_name(func)
register_function_info(func_name, module_name)

@wraps(func)
def sync_wrapper(*args: P.args, **kwds: P.kwargs) -> T:
def sync_wrapper(*args: Params.args, **kwds: Params.kwargs) -> R:
start_time = time.time()
caller_module = caller_module_var.get()
caller_function = caller_function_var.get()
Expand All @@ -121,23 +142,40 @@ def sync_wrapper(*args: P.args, **kwds: P.kwargs) -> T:
if track_concurrency:
track_start(module=module_name, function=func_name)
result = func(*args, **kwds)
track_result_ok(
start_time,
function=func_name,
module=module_name,
caller_module=caller_module,
caller_function=caller_function,
)
if record_error_if and record_error_if(result):
track_result_error(
start_time,
function=func_name,
module=module_name,
caller_module=caller_module,
caller_function=caller_function,
)
else:
track_result_ok(
start_time,
function=func_name,
module=module_name,
caller_module=caller_module,
caller_function=caller_function,
)

except Exception as exception:
result = exception.__class__.__name__
track_result_error(
start_time,
function=func_name,
module=module_name,
caller_module=caller_module,
caller_function=caller_function,
)
if record_success_if and record_success_if(exception):
track_result_ok(
start_time,
function=func_name,
module=module_name,
caller_module=caller_module,
caller_function=caller_function,
)
else:
track_result_error(
start_time,
function=func_name,
module=module_name,
caller_module=caller_module,
caller_function=caller_function,
)
# Reraise exception
raise exception

Expand All @@ -152,15 +190,17 @@ def sync_wrapper(*args: P.args, **kwds: P.kwargs) -> T:
sync_wrapper.__doc__ = append_docs_to_docstring(func, func_name, module_name)
return sync_wrapper

def async_decorator(func: Callable[P, Awaitable[T]]) -> Callable[P, Awaitable[T]]:
def async_decorator(
func: Callable[Params, Awaitable[R]]
) -> Callable[Params, Awaitable[R]]:
"""Helper for decorating async functions, to track calls and duration."""

module_name = get_module_name(func)
func_name = get_function_name(func)
register_function_info(func_name, module_name)

@wraps(func)
async def async_wrapper(*args: P.args, **kwds: P.kwargs) -> T:
async def async_wrapper(*args: Params.args, **kwds: Params.kwargs) -> R:
start_time = time.time()
caller_module = caller_module_var.get()
caller_function = caller_function_var.get()
Expand All @@ -173,23 +213,40 @@ async def async_wrapper(*args: P.args, **kwds: P.kwargs) -> T:
if track_concurrency:
track_start(module=module_name, function=func_name)
result = await func(*args, **kwds)
track_result_ok(
start_time,
function=func_name,
module=module_name,
caller_module=caller_module,
caller_function=caller_function,
)
if record_error_if and record_error_if(result):
track_result_error(
start_time,
function=func_name,
module=module_name,
caller_module=caller_module,
caller_function=caller_function,
)
else:
track_result_ok(
start_time,
function=func_name,
module=module_name,
caller_module=caller_module,
caller_function=caller_function,
)

except Exception as exception:
result = exception.__class__.__name__
track_result_error(
start_time,
function=func_name,
module=module_name,
caller_module=caller_module,
caller_function=caller_function,
)
if record_success_if and record_success_if(exception):
track_result_ok(
start_time,
function=func_name,
module=module_name,
caller_module=caller_module,
caller_function=caller_function,
)
else:
track_result_error(
start_time,
function=func_name,
module=module_name,
caller_module=caller_module,
caller_function=caller_function,
)
# Reraise exception
raise exception

Expand All @@ -204,7 +261,7 @@ async def async_wrapper(*args: P.args, **kwds: P.kwargs) -> T:
async_wrapper.__doc__ = append_docs_to_docstring(func, func_name, module_name)
return async_wrapper

def pick_decorator(func: Callable) -> Callable:
def pick_decorator(func):
"""Pick the correct decorator based on the function type."""
if inspect.iscoroutinefunction(func):
return async_decorator(func)
Expand Down
Empty file added src/autometrics/py.typed
Empty file.
Loading