├── deczoo ├── py.typed ├── _base_notifier.py ├── __init__.py ├── _utils.py └── decorators.py ├── tests ├── __init__.py ├── utils_test │ ├── __init__.py │ ├── empty_shape_error_test.py │ ├── get_free_memory_test.py │ ├── check_parens_test.py │ └── has_shape_test.py ├── conftest.py ├── _base_notifier_test.py └── decorators_test │ ├── raise_if_test.py │ ├── chime_on_end_test.py │ ├── check_args_test.py │ ├── catch_test.py │ ├── notify_on_end_test.py │ ├── timeout_test.py │ ├── memory_limit_test.py │ ├── retry_test.py │ ├── call_counter_test.py │ ├── log_test.py │ ├── shape_tracker_test.py │ └── multi_shape_tracker_test.py ├── docs ├── img │ ├── confused.gif │ ├── deczoo-logo.png │ ├── deeper-meme.jpg │ ├── coverage.svg │ └── interrogate-shield.svg ├── api │ ├── decorators.md │ └── utils.md ├── quickstart.md ├── index.md ├── decorators │ ├── intro.md │ └── advanced.md └── contribute.md ├── CONTRIBUTING.md ├── .github ├── dependabot.yaml └── workflows │ ├── broken-links.yaml │ ├── check-typos.yaml │ ├── pre-commit-update.yaml │ ├── pull-request.yaml │ └── deploy-docs.yaml ├── LICENSE ├── Makefile ├── .pre-commit-config.yaml ├── pyproject.toml ├── .gitignore ├── mkdocs.yml └── README.md /deczoo/py.typed: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /tests/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /tests/utils_test/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /docs/img/confused.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/FBruzzesi/deczoo/HEAD/docs/img/confused.gif -------------------------------------------------------------------------------- /docs/img/deczoo-logo.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/FBruzzesi/deczoo/HEAD/docs/img/deczoo-logo.png -------------------------------------------------------------------------------- /docs/img/deeper-meme.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/FBruzzesi/deczoo/HEAD/docs/img/deeper-meme.jpg -------------------------------------------------------------------------------- /CONTRIBUTING.md: -------------------------------------------------------------------------------- 1 | # Contributing 2 | 3 | Please read the [Contributing guidelines](https://fbruzzesi.github.io/deczoo/contribute/) in the documentation site. 4 | -------------------------------------------------------------------------------- /docs/api/decorators.md: -------------------------------------------------------------------------------- 1 | # Available Decorators 2 | 3 | ::: deczoo.decorators 4 | options: 5 | show_root_full_path: false 6 | show_root_heading: false 7 | -------------------------------------------------------------------------------- /.github/dependabot.yaml: -------------------------------------------------------------------------------- 1 | version: 2 2 | updates: 3 | # Maintain dependencies for GitHub Actions 4 | - package-ecosystem: "github-actions" 5 | directory: "/" 6 | schedule: 7 | interval: "monthly" 8 | -------------------------------------------------------------------------------- /docs/api/utils.md: -------------------------------------------------------------------------------- 1 | # Utility functions 2 | 3 | ::: deczoo._utils.check_parens 4 | options: 5 | show_root_full_path: false 6 | show_root_heading: true 7 | 8 | ::: deczoo._utils._get_free_memory 9 | options: 10 | show_root_full_path: false 11 | show_root_heading: true 12 | -------------------------------------------------------------------------------- /tests/utils_test/empty_shape_error_test.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | 3 | from deczoo._utils import EmptyShapeError 4 | 5 | 6 | def test_empty_shape_error(): 7 | """Test EmptyShapeError class""" 8 | 9 | msg = "Test message" 10 | with pytest.raises(EmptyShapeError) as excinfo: 11 | raise EmptyShapeError(msg) 12 | 13 | assert str(excinfo.value) == msg 14 | -------------------------------------------------------------------------------- /.github/workflows/broken-links.yaml: -------------------------------------------------------------------------------- 1 | name: Check links 2 | 3 | on: 4 | workflow_dispatch: 5 | pull_request: 6 | branches: 7 | - main 8 | 9 | jobs: 10 | 11 | check-links: 12 | runs-on: ubuntu-latest 13 | steps: 14 | - name: Checkout source code 15 | uses: actions/checkout@v5 16 | - name: Markup Link Checker (mlc) 17 | uses: becheran/mlc@v1.0.0 18 | # https://github.com/becheran/mlc 19 | -------------------------------------------------------------------------------- /.github/workflows/check-typos.yaml: -------------------------------------------------------------------------------- 1 | name: Check spelling typos 2 | 3 | on: 4 | workflow_dispatch: 5 | pull_request: 6 | branches: 7 | - main 8 | 9 | jobs: 10 | 11 | run-typos: 12 | runs-on: ubuntu-latest 13 | steps: 14 | - name: Checkout source code 15 | uses: actions/checkout@v5 16 | 17 | - name: Check spelling 18 | uses: crate-ci/typos@master 19 | with: 20 | files: . 21 | -------------------------------------------------------------------------------- /tests/utils_test/get_free_memory_test.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | import pytest 4 | 5 | from deczoo._utils import _get_free_memory 6 | 7 | 8 | @pytest.mark.skipif(os.name != "posix", reason="This test runs only on Unix-based systems") 9 | def test_get_free_memory(): 10 | """ 11 | Tests _get_free_memory function 12 | 13 | The @pytest.mark.skipif(...) decorator skips the test if the operating system is not 14 | a Unix-based system. 15 | """ 16 | assert isinstance(_get_free_memory(), int) 17 | assert _get_free_memory() >= 0 18 | -------------------------------------------------------------------------------- /tests/conftest.py: -------------------------------------------------------------------------------- 1 | import time 2 | 3 | import pytest 4 | 5 | 6 | @pytest.fixture(scope="module") 7 | def base_add(): 8 | """Fixture for base function to decorate""" 9 | 10 | def _add(a, b): 11 | """Adding a and b""" 12 | return a + b 13 | 14 | return _add 15 | 16 | 17 | @pytest.fixture(scope="module") 18 | def sleepy_add(): 19 | """Fixture for base function to decorate""" 20 | 21 | def _add(a, b): 22 | """Adds a and b after sleeping for 2 seconds""" 23 | time.sleep(2) 24 | return a + b 25 | 26 | return _add 27 | -------------------------------------------------------------------------------- /tests/_base_notifier_test.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | 3 | from deczoo._base_notifier import BaseNotifier 4 | 5 | 6 | def test_base_notifier(): 7 | """Tests that BaseNotifier is abstract""" 8 | with pytest.raises(TypeError): 9 | BaseNotifier() 10 | 11 | class TestNotifier(BaseNotifier): 12 | """Concrete implementation of BaseNotifier""" 13 | 14 | def notify(self, *args, **kwargs) -> None: 15 | """Method used to notify""" 16 | print("Notified") 17 | 18 | assert TestNotifier() 19 | assert TestNotifier().notify() is None 20 | -------------------------------------------------------------------------------- /tests/decorators_test/raise_if_test.py: -------------------------------------------------------------------------------- 1 | from contextlib import nullcontext as does_not_raise 2 | 3 | import pytest 4 | 5 | from deczoo import raise_if 6 | 7 | 8 | @pytest.mark.parametrize( 9 | "condition, context", 10 | [ 11 | (lambda: True, pytest.raises(ValueError)), 12 | (lambda: False, does_not_raise()), 13 | ], 14 | ) 15 | def raise_if_test(condition, context): 16 | """Test raise_if decorator.""" 17 | err_msg = "Test exception" 18 | 19 | @raise_if(condition, exception=ValueError, message=err_msg) 20 | def dummy_function(): 21 | pass 22 | 23 | with context as exc_info: 24 | dummy_function() 25 | 26 | if exc_info: 27 | assert err_msg in str(exc_info.value) 28 | -------------------------------------------------------------------------------- /deczoo/_base_notifier.py: -------------------------------------------------------------------------------- 1 | import sys 2 | from abc import ABCMeta, abstractmethod 3 | 4 | if sys.version_info >= (3, 10): 5 | from typing import ParamSpec 6 | else: 7 | from typing_extensions import ParamSpec 8 | 9 | if sys.version_info >= (3, 11): 10 | from typing import Self 11 | else: 12 | from typing_extensions import Self 13 | 14 | PS = ParamSpec("PS") 15 | 16 | 17 | class BaseNotifier(metaclass=ABCMeta): 18 | """Abstract base class to create a notifier to use in `notify_on_end` decorator. 19 | 20 | The class should have a `.notify()` method which gets called after the decorated function has finished running. 21 | """ 22 | 23 | @abstractmethod 24 | def notify(self: Self, *args: PS.args, **kwargs: PS.kwargs) -> None: 25 | """Method used to notify""" 26 | ... 27 | -------------------------------------------------------------------------------- /tests/decorators_test/chime_on_end_test.py: -------------------------------------------------------------------------------- 1 | from unittest.mock import patch 2 | 3 | import pytest 4 | 5 | from deczoo import chime_on_end 6 | 7 | 8 | def test_success(base_add): 9 | """Tests that chime.success is called when function finishes successfully.""" 10 | 11 | with patch("chime.success") as mock_success: 12 | _ = chime_on_end(base_add)(1, 2) 13 | mock_success.assert_called_once() 14 | 15 | 16 | def test_error(): 17 | """Tests that chime.error is called when function raises an error.""" 18 | 19 | with patch("chime.error") as mock_error: 20 | 21 | @chime_on_end 22 | def mock_func(x): 23 | raise ValueError("Something went wrong") 24 | 25 | with pytest.raises(ValueError): 26 | mock_func(10) 27 | 28 | mock_error.assert_called_once 29 | -------------------------------------------------------------------------------- /deczoo/__init__.py: -------------------------------------------------------------------------------- 1 | from importlib import metadata 2 | 3 | from deczoo._base_notifier import BaseNotifier 4 | from deczoo._utils import check_parens 5 | from deczoo.decorators import ( 6 | call_counter, 7 | catch, 8 | check_args, 9 | chime_on_end, 10 | log, 11 | memory_limit, 12 | multi_shape_tracker, 13 | notify_on_end, 14 | raise_if, 15 | retry, 16 | shape_tracker, 17 | timeout, 18 | timer, 19 | ) 20 | 21 | __title__ = __name__ 22 | __version__ = metadata.version(__title__) 23 | 24 | __all__ = ( 25 | "check_parens", 26 | "BaseNotifier", 27 | "call_counter", 28 | "catch", 29 | "check_args", 30 | "chime_on_end", 31 | "log", 32 | "timer", 33 | "memory_limit", 34 | "notify_on_end", 35 | "shape_tracker", 36 | "multi_shape_tracker", 37 | "raise_if", 38 | "retry", 39 | "timeout", 40 | ) 41 | -------------------------------------------------------------------------------- /tests/utils_test/check_parens_test.py: -------------------------------------------------------------------------------- 1 | from deczoo._utils import check_parens 2 | 3 | 4 | def test_check_parens(capsys): 5 | """ 6 | Tests that once a decorator is itself decorated with check_parens, it can be 7 | called with or without parens. 8 | """ 9 | 10 | @check_parens 11 | def decorator(func, arg1="default1", arg2="default2"): 12 | def wrapper(*args, **kwargs): 13 | print(arg1, arg2) 14 | return func(*args, **kwargs) 15 | 16 | return wrapper 17 | 18 | @decorator 19 | def _with_default(): 20 | return "default" 21 | 22 | @decorator(arg1="custom1", arg2="custom2") 23 | def _with_custom(): 24 | return "custom" 25 | 26 | assert _with_default() == "default" 27 | assert _with_custom() == "custom" 28 | 29 | all_args = ("default1", "default2", "custom1", "custom2") 30 | sys_out = capsys.readouterr().out 31 | 32 | assert all(a in sys_out for a in all_args) 33 | -------------------------------------------------------------------------------- /docs/img/coverage.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | coverage 17 | coverage 18 | 99% 19 | 99% 20 | 21 | 22 | -------------------------------------------------------------------------------- /.github/workflows/pre-commit-update.yaml: -------------------------------------------------------------------------------- 1 | name: Pre-commit auto-update 2 | 3 | on: 4 | workflow_dispatch: 5 | schedule: 6 | - cron: "0 0 1 * *" # Every 1st of the month at 00:00 UTC 7 | 8 | permissions: write-all 9 | 10 | jobs: 11 | auto-update: 12 | runs-on: ubuntu-latest 13 | steps: 14 | - name: Checkout source code 15 | uses: actions/checkout@v5 16 | 17 | - name: Set up Python 18 | uses: actions/setup-python@v6 19 | with: 20 | python-version: "3.10" 21 | 22 | - name: pre-commit install autoupdate 23 | run: | 24 | pip install pre-commit 25 | pre-commit autoupdate 26 | 27 | - name: Commit and push changes 28 | uses: peter-evans/create-pull-request@v7 29 | with: 30 | branch: update-pre-commit-hooks 31 | title: 'Update pre-commit hooks' 32 | commit-message: 'Update pre-commit hooks' 33 | body: | 34 | Update versions of pre-commit hooks to latest versions. 35 | -------------------------------------------------------------------------------- /tests/utils_test/has_shape_test.py: -------------------------------------------------------------------------------- 1 | from typing import Tuple 2 | 3 | import pytest 4 | 5 | from deczoo._utils import SupportShape 6 | 7 | 8 | class TestShape: 9 | """Test class for HasShape with shape attribute""" 10 | 11 | @property 12 | def shape(self) -> Tuple[int, ...]: 13 | return (1, 2) 14 | 15 | 16 | class NoShape: 17 | """Test class for HasShape without shape attribute""" 18 | 19 | pass 20 | 21 | 22 | def test_protocol(): 23 | """Tests that HasShape is protocol and cannot be instantiated""" 24 | with pytest.raises(TypeError): 25 | SupportShape() 26 | 27 | 28 | def test_protocol_implemented(): 29 | """Tests that TestShape is a valid implementation of HasShape""" 30 | assert isinstance(TestShape, SupportShape) 31 | assert hasattr(TestShape, "shape") 32 | 33 | 34 | def test_protocol_not_implemented(): 35 | """Tests that NoShape is not a valid implementation of HasShape""" 36 | 37 | assert not isinstance(NoShape, SupportShape) 38 | assert not hasattr(NoShape, "shape") 39 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2022 Francesco Bruzzesi 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /tests/decorators_test/check_args_test.py: -------------------------------------------------------------------------------- 1 | from contextlib import nullcontext as does_not_raise 2 | 3 | import pytest 4 | 5 | from deczoo import check_args 6 | 7 | 8 | @pytest.mark.parametrize( 9 | "rules, context", 10 | [ 11 | ({"a": lambda t: t > 0, "b": lambda t: t < 10}, does_not_raise()), 12 | ({}, does_not_raise()), 13 | ({"a": True}, pytest.raises(ValueError)), 14 | ({"a": lambda t: t > 0, "b": "test"}, pytest.raises(ValueError)), 15 | ], 16 | ) 17 | def test_params(base_add, rules, context): 18 | """ 19 | Tests that check_args raises an error if invalid parameter is passed. 20 | """ 21 | 22 | with context: 23 | check_args(base_add, **rules) 24 | 25 | 26 | @pytest.mark.parametrize( 27 | "rules, context", 28 | [ 29 | ({"a": lambda t: t > 0, "b": lambda t: t > 0}, does_not_raise()), 30 | ({"a": lambda t: t < 2, "b": lambda t: t > 0}, does_not_raise()), 31 | ({"a": lambda t: t > 0, "b": lambda t: t < 0}, pytest.raises(ValueError)), 32 | ({"a": lambda t: t < 0, "b": lambda t: t > 0}, pytest.raises(ValueError)), 33 | ], 34 | ) 35 | def test_rules(base_add, rules, context): 36 | """ 37 | Tests that check_args applies the rules correctly. 38 | """ 39 | add = check_args(base_add, **rules) 40 | 41 | with context: 42 | add(a=1, b=1) 43 | -------------------------------------------------------------------------------- /Makefile: -------------------------------------------------------------------------------- 1 | init-env: 2 | pip install . --no-cache-dir 3 | 4 | init-dev: 5 | pip install -e ".[all-dev]" --no-cache-dir 6 | pre-commit install 7 | 8 | clean-notebooks: 9 | jupyter nbconvert --ClearOutputPreprocessor.enabled=True --inplace notebooks/*.ipynb 10 | 11 | clean-folders: 12 | rm -rf __pycache__ */__pycache__ */**/__pycache__ \ 13 | .pytest_cache */.pytest_cache */**/.pytest_cache \ 14 | .ruff_cache */.ruff_cache */**/.ruff_cache \ 15 | .mypy_cache */.mypy_cache */**/.mypy_cache \ 16 | site build dist htmlcov .coverage .tox 17 | 18 | lint: 19 | ruff version 20 | ruff check deczoo tests --fix 21 | ruff format deczoo tests 22 | ruff clean 23 | 24 | test: 25 | pytest tests -n auto 26 | 27 | coverage: 28 | rm -rf .coverage 29 | (rm docs/img/coverage.svg) || (echo "No coverage.svg file found") 30 | coverage run -m pytest 31 | coverage report -m 32 | coverage-badge -o docs/img/coverage.svg 33 | 34 | interrogate: 35 | interrogate deczoo tests 36 | 37 | interrogate-badge: 38 | interrogate --generate-badge docs/img/interrogate-shield.svg 39 | 40 | check: interrogate lint test clean-folders 41 | 42 | docs-serve: 43 | mkdocs serve 44 | 45 | docs-deploy: 46 | mkdocs gh-deploy 47 | 48 | pypi-push: 49 | rm -rf dist 50 | hatch build 51 | hatch publish 52 | 53 | get-version : 54 | @echo $(shell grep -m 1 version pyproject.toml | tr -s ' ' | tr -d '"' | tr -d "'" | cut -d' ' -f3) 55 | -------------------------------------------------------------------------------- /.pre-commit-config.yaml: -------------------------------------------------------------------------------- 1 | # See https://pre-commit.com for more information 2 | # See https://pre-commit.com/hooks.html for more hooks 3 | repos: 4 | - repo: https://github.com/pre-commit/pre-commit-hooks 5 | rev: v6.0.0 6 | hooks: 7 | - id: trailing-whitespace 8 | - id: name-tests-test 9 | - id: end-of-file-fixer 10 | - id: requirements-txt-fixer 11 | - id: check-json 12 | - id: check-yaml 13 | - id: check-ast 14 | - id: check-added-large-files 15 | - repo: https://github.com/astral-sh/ruff-pre-commit 16 | rev: v0.13.2 # Ruff version. 17 | hooks: 18 | - id: ruff # Run the linter. 19 | args: [--fix, deczoo, tests] 20 | - id: ruff-format # Run the formatter. 21 | args: [deczoo, tests] 22 | - repo: https://github.com/Lucas-C/pre-commit-hooks-bandit 23 | rev: v1.0.6 24 | hooks: 25 | - id: python-bandit-vulnerability-check 26 | args: [--skip, "B101",--severity-level, medium, --recursive, deczoo] 27 | - repo: https://github.com/econchick/interrogate 28 | rev: 1.7.0 29 | hooks: 30 | - id: interrogate 31 | args: [-vv, --ignore-nested-functions, --ignore-module, --ignore-init-method, --ignore-private, --ignore-magic, --ignore-property-decorators, --fail-under=90, deczoo, tests] 32 | - repo: https://github.com/pre-commit/pygrep-hooks 33 | rev: v1.10.0 34 | hooks: 35 | - id: python-no-eval 36 | -------------------------------------------------------------------------------- /tests/decorators_test/catch_test.py: -------------------------------------------------------------------------------- 1 | from contextlib import nullcontext as does_not_raise 2 | 3 | import pytest 4 | 5 | from deczoo import catch 6 | 7 | 8 | @pytest.mark.parametrize( 9 | "value, context", 10 | [ 11 | (1.1, pytest.raises(TypeError)), 12 | ("a", pytest.raises(TypeError)), 13 | ((1, 2), pytest.raises(TypeError)), 14 | (print, does_not_raise()), 15 | ], 16 | ) 17 | def test_params(base_add, value, context): 18 | """ 19 | Tests that catch parse its params. 20 | """ 21 | 22 | with context: 23 | catch(base_add, logging_fn=value) 24 | 25 | 26 | @pytest.mark.parametrize( 27 | "b, return_on_exception, expected", 28 | [(2, -999, 3), ("a", -999, -999), ("a", 0, 0)], 29 | ) 30 | def test_return(base_add, b, return_on_exception, expected): 31 | """Tests that catch returns return_on_exception value""" 32 | add = catch(base_add, return_on_exception=return_on_exception) 33 | 34 | assert add(a=1, b=b) == expected 35 | 36 | 37 | @pytest.mark.parametrize( 38 | "raise_on_exception, context", 39 | [ 40 | (ValueError, pytest.raises(ValueError)), 41 | (Exception, pytest.raises(Exception)), 42 | (None, pytest.raises(Exception)), 43 | ], 44 | ) 45 | def test_raise(base_add, raise_on_exception, context): 46 | """Tests that catch raises the raise_on_exception exception""" 47 | add = catch(base_add, raise_on_exception=raise_on_exception) 48 | 49 | with context: 50 | add(a=1, b="a") 51 | -------------------------------------------------------------------------------- /tests/decorators_test/notify_on_end_test.py: -------------------------------------------------------------------------------- 1 | from contextlib import nullcontext as does_not_raise 2 | 3 | import pytest 4 | 5 | from deczoo import notify_on_end 6 | from deczoo._base_notifier import BaseNotifier 7 | 8 | 9 | class TestNotifier(BaseNotifier): 10 | """ 11 | Test notifier class, concrete implementation of BaseNotifier used for testing 12 | """ 13 | 14 | def notify(self, *args, **kwargs): 15 | """Method used to notify""" 16 | print("Notified") 17 | 18 | 19 | class TestNotNotifier: 20 | """Test class that is not a notifier""" 21 | 22 | def notify(self, *args, **kwargs): 23 | """Method used to notify""" 24 | ... 25 | 26 | 27 | @pytest.mark.parametrize( 28 | "notifier, context", 29 | [(TestNotifier(), does_not_raise()), (TestNotNotifier(), pytest.raises(TypeError))], 30 | ) 31 | def test_params(base_add, notifier, context): 32 | """Tests that notify_on_end raises an error if invalid parameter is passed.""" 33 | 34 | with context: 35 | notify_on_end(base_add, notifier=notifier) 36 | 37 | 38 | @pytest.mark.parametrize( 39 | "a, b, exception", 40 | [(1, 2, does_not_raise()), (1, "a", pytest.raises(Exception))], 41 | ) 42 | def test_notify(base_add, capsys, a, b, exception): 43 | """Tests that notify method get called""" 44 | 45 | add = notify_on_end(base_add, notifier=TestNotifier()) 46 | 47 | with exception: 48 | add(a=a, b=b) 49 | 50 | sys_out = capsys.readouterr().out 51 | assert "Notified" in sys_out 52 | -------------------------------------------------------------------------------- /tests/decorators_test/timeout_test.py: -------------------------------------------------------------------------------- 1 | from contextlib import nullcontext as does_not_raise 2 | 3 | import pytest 4 | 5 | from deczoo import timeout 6 | 7 | 8 | @pytest.mark.parametrize( 9 | "arg_name, value, context", 10 | [ 11 | ("time_limit", -1, pytest.raises(ValueError)), 12 | ("time_limit", 1.0, pytest.raises(ValueError)), 13 | ("time_limit", "a", pytest.raises(ValueError)), 14 | ("time_limit", (1, 2), pytest.raises(ValueError)), 15 | ("signum", "a", pytest.raises(TypeError)), 16 | ("signum", (1, 2), pytest.raises(TypeError)), 17 | ("signal_handler", "a", pytest.raises(TypeError)), 18 | ("signal_handler", (1, 2), pytest.raises(TypeError)), 19 | ("time_limit", 1, does_not_raise()), 20 | ("signum", 1, does_not_raise()), 21 | ("signal_handler", lambda x, y: (x, y), does_not_raise()), 22 | ], 23 | ) 24 | def test_params(base_add, arg_name, value, context): 25 | """ 26 | Tests that timeout raises an error if invalid parameter is passed. 27 | """ 28 | 29 | with context: 30 | if arg_name != "time_limit": 31 | timeout(base_add, time_limit=1, **{arg_name: value}) 32 | else: 33 | timeout(base_add, **{arg_name: value}) 34 | 35 | 36 | @pytest.mark.parametrize( 37 | "b, time_limit, context", 38 | [ 39 | (1, 1, pytest.raises(TimeoutError)), 40 | (1, 2, pytest.raises(TimeoutError)), 41 | ( 42 | "a", 43 | 1, 44 | pytest.raises(TimeoutError), 45 | ), # TimeoutError is raised before TypeError 46 | (1, 3, does_not_raise()), 47 | ("a", 3, pytest.raises(TypeError)), 48 | ], 49 | ) 50 | def test_out_of_limit(sleepy_add, b, time_limit, context): 51 | """Tests that if the function doesn't make it in time TimeoutError is raised""" 52 | 53 | with context: 54 | _ = timeout(sleepy_add, time_limit=time_limit)(1, b=b) 55 | -------------------------------------------------------------------------------- /tests/decorators_test/memory_limit_test.py: -------------------------------------------------------------------------------- 1 | from contextlib import nullcontext as does_not_raise 2 | 3 | import pytest 4 | 5 | from deczoo import memory_limit 6 | 7 | 8 | @pytest.mark.parametrize( 9 | "arg_name, value, context", 10 | [ 11 | ("percentage", 2, pytest.raises(TypeError)), 12 | ("percentage", "a", pytest.raises(TypeError)), 13 | ("percentage", (1, 2), pytest.raises(TypeError)), 14 | ("logging_fn", 2, pytest.raises(TypeError)), 15 | ("logging_fn", "a", pytest.raises(TypeError)), 16 | ("logging_fn", (1, 2), pytest.raises(TypeError)), 17 | ("percentage", 2.0, pytest.raises(ValueError)), 18 | ("percentage", -1.0, pytest.raises(ValueError)), 19 | ("percentage", 0.0, does_not_raise()), 20 | ("percentage", 0.5, does_not_raise()), 21 | ("percentage", 1.0, does_not_raise()), 22 | ("logging_fn", print, does_not_raise()), 23 | ], 24 | ) 25 | def test_params(base_add, arg_name, value, context): 26 | """ 27 | Tests that memory_limit raises an error if invalid parameter is passed. 28 | """ 29 | 30 | with context: 31 | memory_limit(base_add, **{arg_name: value}) 32 | 33 | 34 | @pytest.mark.parametrize( 35 | "percentage, context", 36 | [ 37 | (0.01, pytest.raises(MemoryError)), 38 | (0.02, pytest.raises(MemoryError)), 39 | (0.95, does_not_raise()), 40 | (1.0, does_not_raise()), 41 | ], 42 | ) 43 | def test_memory_limit(capsys, percentage, context): 44 | """Tests that memory limited function raises MemoryError exception""" 45 | 46 | @memory_limit(percentage=percentage, logging_fn=print) 47 | def limited(x): 48 | for i in list(range(10**7)): 49 | _ = 1 + 1 50 | return x 51 | 52 | with context: 53 | limited(42) 54 | 55 | sys_out = capsys.readouterr().out 56 | assert f"Setting memory limit for {limited.__name__} to" in sys_out 57 | -------------------------------------------------------------------------------- /tests/decorators_test/retry_test.py: -------------------------------------------------------------------------------- 1 | from contextlib import nullcontext as does_not_raise 2 | 3 | import pytest 4 | 5 | from deczoo import retry 6 | 7 | 8 | @pytest.mark.parametrize( 9 | "arg_name, value, context", 10 | [ 11 | ("n_tries", 0, pytest.raises(ValueError)), 12 | ("n_tries", 0.5, pytest.raises(ValueError)), 13 | ("delay", "a", pytest.raises(ValueError)), 14 | ("delay", -2, pytest.raises(ValueError)), 15 | ("logging_fn", "a", pytest.raises(TypeError)), 16 | ("logging_fn", (1, 2), pytest.raises(TypeError)), 17 | ("n_tries", 1, does_not_raise()), 18 | ("n_tries", 2, does_not_raise()), 19 | ("delay", 1, does_not_raise()), 20 | ("delay", 1.0, does_not_raise()), 21 | ("logging_fn", print, does_not_raise()), 22 | ], 23 | ) 24 | def test_params(base_add, arg_name, value, context): 25 | """ 26 | Tests that retry raises an error if invalid parameter is passed. 27 | """ 28 | 29 | with context: 30 | retry(base_add, **{arg_name: value}) 31 | 32 | 33 | @pytest.mark.parametrize("n_tries", list(range(2, 5))) 34 | @pytest.mark.parametrize("delay", [0.0, 0.1]) 35 | def test_retry_no_exceptions(base_add, capsys, n_tries, delay): 36 | """Tests that retry decorator works""" 37 | _ = retry(base_add, n_tries=n_tries, delay=delay, logging_fn=print)(a=1, b=2) 38 | assert f"Attempt 1/{n_tries}: Succeeded" in capsys.readouterr().out 39 | 40 | 41 | @pytest.mark.parametrize("n_tries", list(range(2, 5))) 42 | @pytest.mark.parametrize("delay", [0.0, 0.1]) 43 | def test_retry_with_exceptions(base_add, capsys, n_tries, delay): 44 | """Tests that retry decorator retries when exceptions are raised""" 45 | add = retry(base_add, n_tries=n_tries, delay=delay, logging_fn=print) 46 | with pytest.raises(Exception): 47 | add(a=1, b="a") 48 | 49 | sys_out = capsys.readouterr().out 50 | assert all(f"Attempt {x}/{n_tries}: Failed" in sys_out for x in range(1, n_tries + 1)) 51 | -------------------------------------------------------------------------------- /tests/decorators_test/call_counter_test.py: -------------------------------------------------------------------------------- 1 | from contextlib import nullcontext as does_not_raise 2 | 3 | import pytest 4 | 5 | from deczoo import call_counter 6 | 7 | 8 | @pytest.mark.parametrize("arg_name", ["seed", "log_counter", "logging_fn"]) 9 | @pytest.mark.parametrize("value", [1.1, "a", (1, 2)]) 10 | def test_params_raise(base_add, arg_name, value): 11 | """ 12 | Tests that call_counter raises an error if invalid parameter is passed. 13 | """ 14 | 15 | with pytest.raises(TypeError): 16 | call_counter(base_add, **{arg_name: value}) 17 | 18 | 19 | @pytest.mark.parametrize( 20 | "arg_name, value", 21 | [("seed", 1), ("log_counter", True), ("log_counter", False), ("logging_fn", print)], 22 | ) 23 | def test_params_valid(base_add, arg_name, value): 24 | """ 25 | Tests that call_counter doesn't raise an error if valid parameter type is passed. 26 | """ 27 | 28 | with does_not_raise(): 29 | call_counter(base_add, **{arg_name: value}) 30 | 31 | 32 | @pytest.mark.parametrize("seed", list(range(-10, 10, 2))) 33 | @pytest.mark.parametrize("n_calls", list(range(1, 5))) 34 | def test_counting(base_add, seed, n_calls): 35 | """ 36 | Tests that call_counter keeps track of number of times function has been called, 37 | starting from the given seed. 38 | """ 39 | expected = seed + n_calls 40 | add = call_counter(base_add, seed=seed) 41 | 42 | for _ in range(n_calls): 43 | add(1, 2) 44 | 45 | assert add._calls == expected 46 | 47 | 48 | @pytest.mark.parametrize("n_calls", list(range(1, 5))) 49 | def test_sys(capsys, base_add, n_calls): 50 | """ 51 | Tests that call_counter logs the number of times function has been called. 52 | """ 53 | add = call_counter(base_add, seed=0, log_counter=True, logging_fn=print) 54 | 55 | for _ in range(n_calls): 56 | add(1, 2) 57 | 58 | sys_out = capsys.readouterr().out 59 | assert all(f"called {x} times" in sys_out for x in range(1, n_calls + 1)) 60 | -------------------------------------------------------------------------------- /docs/quickstart.md: -------------------------------------------------------------------------------- 1 | # Getting started 2 | 3 | The idea is kind of simple: each function in the library is a (function) decorator with a specific objective in mind. 4 | 5 | ```python title="Example: log decorator" 6 | from deczoo import log 7 | 8 | @log # equivalent to @log(log_time=True, log_args=True, log_error=True, logging_fn=print) 9 | def custom_add(a, b, *args): 10 | """Adds all arguments together""" 11 | return sum([a, b, *args]) 12 | 13 | _ = custom_add(1, 2, 3, 4) 14 | # custom_add args=(a=1, b=2, args=(3, 4)) time=0:00:00.000062 15 | 16 | _ = custom_add(1, "a", 2) 17 | # custom_add args=(a=1, b=a, args=(2,)) time=0:00:00.000064 Failed with error: unsupported 18 | # operand type(s) for +: 'int' and 'str' 19 | ``` 20 | 21 | ```python title="Example: shape_tracker decorator" 22 | from deczoo import shape_tracker 23 | 24 | @shape_tracker(shape_in=True, shape_out=True, shape_delta=True, raise_if_empty=True) 25 | def tracked_vstack(a: np.ndarray, b: np.ndarray) -> np.ndarray: 26 | return np.vstack([a, b]) 27 | 28 | _ = tracked_vstack(np.ones((1, 2)), np.ones((10, 2))) 29 | # Input: `a` has shape (1, 2) 30 | # Output: result has shape (11, 2) 31 | # Shape delta: (-10, 0) 32 | ``` 33 | 34 | ## Features 35 | 36 | The library implements the following decorators: 37 | 38 | - `call_counter`: tracks how many times a function has been called. 39 | - `catch`: wraps a function in a try-except block, returning a custom value, or raising a custom exception. 40 | - `check_args`: checks that function arguments satisfy its "rule". 41 | - `chime_on_end`: notify with chime sound on function end (success or error). 42 | - `log`: tracks function time taken, arguments and errors, such logs can be written to a file. 43 | - `timer`: tracks function time taken. 44 | - `memory_limit`: sets a memory limit while running the function. 45 | - `notify_on_end`: notifies when function finished running with a custom notifier. 46 | - `retry`: wraps a function with a "retry" block. 47 | - `shape_tracker`: tracks the shape of a dataframe/array-like object, in input and/or output. 48 | - `multi_shape_tracker`: tracks the shapes of input(s) and/or output(s) of a function. 49 | - `timeout`: sets a time limit for the function, terminates the process if it hasn't finished within such time limit. 50 | 51 | ## Examples 52 | 53 | Please refer to the [api page](api/decorators.md) to see a basic example for each decorator. 54 | -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [build-system] 2 | requires = ["hatchling"] 3 | build-backend = "hatchling.build" 4 | 5 | [project] 6 | name = "deczoo" 7 | version = "0.6.0" 8 | description = "Zoo for Python decorators" 9 | 10 | license = {file = "LICENSE"} 11 | readme = "README.md" 12 | requires-python = ">=3.8" 13 | authors = [{name = "Francesco Bruzzesi"}] 14 | 15 | dependencies = [ 16 | "typing-extensions>=4.4.0; python_version < '3.12'", 17 | ] 18 | 19 | classifiers = [ 20 | "Programming Language :: Python :: 3", 21 | "Programming Language :: Python :: 3.8", 22 | "Programming Language :: Python :: 3.9", 23 | "Programming Language :: Python :: 3.10", 24 | "Programming Language :: Python :: 3.11", 25 | "Programming Language :: Python :: 3.12", 26 | "License :: OSI Approved :: MIT License" 27 | ] 28 | 29 | [project.urls] 30 | documentation = "https://fbruzzesi.github.io/deczoo/" 31 | repository = "https://github.com/fbruzzesi/deczoo" 32 | issue-tracker = "https://github.com/fbruzzesi/deczoo/issues" 33 | 34 | 35 | [project.optional-dependencies] 36 | chime = ["chime"] 37 | rich = ["rich>=12.0.0"] 38 | 39 | dev = [ 40 | "pre-commit>=2.21.0", 41 | "hatch" 42 | ] 43 | 44 | lint = [ 45 | "ruff>=0.1.6" 46 | ] 47 | 48 | test = [ 49 | "interrogate>=1.5.0", 50 | "pytest>=7.2.0", 51 | "pytest-xdist>=3.2.1", 52 | "coverage>=7.2.1", 53 | "coverage-badge>=1.1.0", 54 | "numpy", 55 | ] 56 | 57 | docs = [ 58 | "mkdocs>=1.4.2", 59 | "mkdocs-material>=9.1.2", 60 | "mkdocstrings[python]>=0.20.0", 61 | ] 62 | 63 | all = ["deczoo[chime,rich]"] 64 | all-dev = ["deczoo[chime,rich,dev,lint,test,docs]"] 65 | 66 | [tool.hatch.build.targets.sdist] 67 | only-include = ["deczoo"] 68 | 69 | [tool.hatch.build.targets.wheel] 70 | packages = ["deczoo"] 71 | 72 | [tool.ruff] 73 | line-length = 120 74 | extend-select = ["I"] 75 | ignore = [ 76 | "E731", # do not assign a `lambda` expression, use a `def` 77 | ] 78 | 79 | [tool.ruff.lint.pydocstyle] 80 | convention = "google" 81 | 82 | [tool.interrogate] 83 | ignore-nested-functions = true 84 | ignore-module = true 85 | ignore-init-method = true 86 | ignore-private = true 87 | ignore-magic = true 88 | ignore-property-decorators = true 89 | fail-under = 95 90 | verbose = 2 # 0 (minimal output), 1 (-v), 2 (-vv) 91 | 92 | [tool.mypy] 93 | ignore_missing_imports = true 94 | 95 | [tool.coverage.run] 96 | source = ["deczoo/"] 97 | -------------------------------------------------------------------------------- /tests/decorators_test/log_test.py: -------------------------------------------------------------------------------- 1 | from contextlib import nullcontext as does_not_raise 2 | 3 | import pytest 4 | 5 | from deczoo import log 6 | 7 | 8 | @pytest.mark.parametrize( 9 | "arg_name, value, context", 10 | [ 11 | ("log_time", 1.1, pytest.raises(TypeError)), 12 | ("log_args", "a", pytest.raises(TypeError)), 13 | ("log_error", (1, 2), pytest.raises(TypeError)), 14 | ("log_file", 1.1, pytest.raises(TypeError)), 15 | ("logging_fn", {}, pytest.raises(TypeError)), 16 | ("log_time", True, does_not_raise()), 17 | ("log_args", False, does_not_raise()), 18 | ("log_error", False, does_not_raise()), 19 | ("log_file", "test.txt", does_not_raise()), 20 | ("logging_fn", print, does_not_raise()), 21 | ], 22 | ) 23 | def test_params(base_add, arg_name, value, context): 24 | """Tests that log raises an error if invalid parameter is passed.""" 25 | 26 | with context: 27 | log(base_add, **{arg_name: value}) 28 | 29 | 30 | @pytest.mark.parametrize( 31 | "log_time, log_args, expected", 32 | [ 33 | (True, False, "add time="), 34 | (False, True, "add args=(a=1, b=2)"), 35 | (True, True, "add args=(a=1, b=2) time="), 36 | ], 37 | ) 38 | def test_log_no_exceptions(base_add, capsys, log_time, log_args, expected): 39 | """Tests that log decorator works""" 40 | 41 | add = log( 42 | base_add, 43 | log_time=log_time, 44 | log_args=log_args, 45 | logging_fn=print, 46 | ) 47 | 48 | add(a=1, b=2) 49 | 50 | sys_out = capsys.readouterr().out 51 | assert expected in sys_out 52 | 53 | 54 | def test_log_with_exceptions(base_add, capsys): 55 | """Tests that log decorator logs exceptions""" 56 | 57 | add = log(base_add, log_time=False, log_args=False, log_error=True, logging_fn=print) 58 | 59 | with pytest.raises(Exception): 60 | add(a=1, b="a") 61 | 62 | sys_out = capsys.readouterr().out 63 | assert "add Failed with error: " in sys_out 64 | 65 | 66 | def test_log_file(base_add, tmp_path): 67 | """Tests that log decorator logs to file""" 68 | 69 | log_file = tmp_path / "log.txt" 70 | 71 | add = log( 72 | base_add, 73 | log_time=False, 74 | log_args=True, 75 | log_file=log_file, 76 | logging_fn=print, 77 | ) 78 | 79 | add(a=1, b=2) 80 | with open(log_file) as f: 81 | assert "add args=(a=1, b=2)" in f.read() 82 | -------------------------------------------------------------------------------- /.github/workflows/pull-request.yaml: -------------------------------------------------------------------------------- 1 | name: PR Checks 2 | 3 | on: 4 | pull_request: 5 | branches: 6 | - main 7 | 8 | jobs: 9 | 10 | interrogate: 11 | runs-on: ubuntu-latest 12 | steps: 13 | - name: Checkout source code 14 | uses: actions/checkout@v5 15 | - name: Set up Python 16 | uses: actions/setup-python@v6 17 | with: 18 | python-version: "3.10" 19 | - name: Install dependencies and run interrogate 20 | run: | 21 | python -m pip install --upgrade pip --no-cache-dir 22 | python -m pip install interrogate==1.5.0 --no-cache-dir 23 | make interrogate 24 | 25 | lint: 26 | needs: [interrogate] 27 | runs-on: ubuntu-latest 28 | steps: 29 | - name: Checkout source code 30 | uses: actions/checkout@v5 31 | - name: Set up Python 32 | uses: actions/setup-python@v6 33 | with: 34 | python-version: "3.10" 35 | - name: Install dependencies and run linter 36 | run: | 37 | python -m pip install --upgrade pip 38 | python -m pip install ."[lint]" --no-cache-dir 39 | make lint 40 | 41 | test: 42 | needs: [lint] 43 | strategy: 44 | matrix: 45 | os: [ubuntu-latest 46 | #, macos-latest 47 | #, windows-latest 48 | ] 49 | python-version: ["3.8", "3.9", "3.10", "3.11", "3.12"] 50 | runs-on: ${{ matrix.os }} 51 | steps: 52 | - name: Checkout source code 53 | uses: actions/checkout@v5 54 | - name: Set up Python ${{ matrix.python-version }} 55 | uses: actions/setup-python@v6 56 | with: 57 | python-version: ${{ matrix.python-version }} 58 | - name: Install dependencies and run tests 59 | run: | 60 | python -m pip install --upgrade pip --no-cache-dir 61 | python -m pip install pytest==7.2.0 pytest-xdist==3.2.1 numpy --no-cache-dir 62 | python -m pip install ."[chime]" --no-cache-dir 63 | - name: Test 64 | run: make test 65 | 66 | doc-build: 67 | needs: [test] 68 | runs-on: ubuntu-latest 69 | steps: 70 | - name: Checkout source code 71 | uses: actions/checkout@v5 72 | - name: Set up Python 73 | uses: actions/setup-python@v6 74 | with: 75 | python-version: "3.10" 76 | - name: Install dependencies and check docs can build 77 | run: | 78 | python -m pip install --upgrade pip --no-cache-dir 79 | python -m pip install ."[docs]" --no-cache-dir 80 | mkdocs build -v -s 81 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | .ruff_cache/ 2 | 3 | # Notebooks 4 | notebooks/ 5 | *.ipynb 6 | 7 | # Files 8 | *.json 9 | *.csv 10 | *.xlsx 11 | *.pickle 12 | *.pkl 13 | *.parquet 14 | 15 | 16 | # Byte-compiled / optimized / DLL files 17 | __pycache__/ 18 | *.py[cod] 19 | *$py.class 20 | 21 | # C extensions 22 | *.so 23 | 24 | # Distribution / packaging 25 | .Python 26 | build/ 27 | develop-eggs/ 28 | dist/ 29 | downloads/ 30 | eggs/ 31 | .eggs/ 32 | lib/ 33 | lib64/ 34 | parts/ 35 | sdist/ 36 | var/ 37 | wheels/ 38 | pip-wheel-metadata/ 39 | share/python-wheels/ 40 | *.egg-info/ 41 | .installed.cfg 42 | *.egg 43 | MANIFEST 44 | 45 | # PyInstaller 46 | # Usually these files are written by a python script from a template 47 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 48 | *.manifest 49 | *.spec 50 | 51 | # Installer logs 52 | pip-log.txt 53 | pip-delete-this-directory.txt 54 | 55 | # Unit test / coverage reports 56 | htmlcov/ 57 | .tox/ 58 | .nox/ 59 | .coverage 60 | .coverage.* 61 | .cache 62 | nosetests.xml 63 | coverage.xml 64 | *.cover 65 | *.py,cover 66 | .hypothesis/ 67 | .pytest_cache/ 68 | 69 | # Translations 70 | *.mo 71 | *.pot 72 | 73 | # Django stuff: 74 | *.log 75 | local_settings.py 76 | db.sqlite3 77 | db.sqlite3-journal 78 | 79 | # Flask stuff: 80 | instance/ 81 | .webassets-cache 82 | 83 | # Scrapy stuff: 84 | .scrapy 85 | 86 | # Sphinx documentation 87 | docs/_build/ 88 | 89 | # PyBuilder 90 | target/ 91 | 92 | # Jupyter Notebook 93 | .ipynb_checkpoints 94 | 95 | # IPython 96 | profile_default/ 97 | ipython_config.py 98 | 99 | # pyenv 100 | .python-version 101 | 102 | # pipenv 103 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 104 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 105 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 106 | # install all needed dependencies. 107 | #Pipfile.lock 108 | 109 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow 110 | __pypackages__/ 111 | 112 | # Celery stuff 113 | celerybeat-schedule 114 | celerybeat.pid 115 | 116 | # SageMath parsed files 117 | *.sage.py 118 | 119 | # Environments 120 | .env 121 | .venv 122 | env/ 123 | venv/ 124 | ENV/ 125 | env.bak/ 126 | venv.bak/ 127 | 128 | # Spyder project settings 129 | .spyderproject 130 | .spyproject 131 | 132 | # Rope project settings 133 | .ropeproject 134 | 135 | # mkdocs documentation 136 | /site 137 | 138 | # mypy 139 | .mypy_cache/ 140 | .dmypy.json 141 | dmypy.json 142 | 143 | # Pyre type checker 144 | .pyre/ 145 | -------------------------------------------------------------------------------- /docs/index.md: -------------------------------------------------------------------------------- 1 | 2 | 3 | ![](https://img.shields.io/github/license/FBruzzesi/deczoo) 4 | 5 | [![Code style: black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black) 6 | 7 | 8 | 9 | 10 | # Deczoo 11 | 12 | > A zoo for decorators 13 | 14 | There are many great decorators out there that we use everyday. Why don't collect a few of them? 15 | 16 | I found myself implementing over and over some common decorators in different projects. 17 | The hope is to gather them here and use this codebase. 18 | 19 | --- 20 | 21 | [Documentation](https://fbruzzesi.github.io/deczoo/) | [Source Code](https://github.com/fbruzzesi/deczoo) 22 | 23 | --- 24 | 25 | ## Alpha Notice 26 | 27 | This codebase is experimental and is working for my use cases. It is very probable that there are cases not covered and for which it breaks (badly). If you find them, please feel free to open an issue in the [issue page](https://github.com/FBruzzesi/deczoo/issues) of the repo. 28 | 29 | ## What is a decorator? 30 | 31 | In short a python decorator is a way to modify or enhance the behavior of a function or a class without actually modifying the source code of the function or class. 32 | 33 | Decorators are implemented as functions (or classes) that take a function or a class as input and return a new function or class that has some additional functionality. 34 | 35 | To have a more in-depth explanation you can check the [next section](decorators/intro.md). 36 | 37 | ## Installation 38 | 39 | **deczoo** is published as a Python package on [pypi](https://pypi.org/), and it can be installed with pip, directly from source using git, or with a local clone: 40 | 41 | === "pip (pypi)" 42 | 43 | ```bash 44 | python -m pip install deczoo 45 | ``` 46 | 47 | === "source/git" 48 | 49 | ```bash 50 | python -m pip install git+https://github.com/FBruzzesi/deczoo.git 51 | ``` 52 | 53 | === "local clone" 54 | 55 | ```bash 56 | git clone https://github.com/FBruzzesi/deczoo.git 57 | cd deczoo 58 | python -m pip install . 59 | ``` 60 | 61 | ### Dependencies 62 | 63 | As of now, the library has no additional required dependencies, however: 64 | 65 | - some functionalities works only on UNIX systems (`@memory_limit` and `@timeout`) 66 | - to use some decorators you may need to install additional dependencies (e.g. install [`chime`](https://github.com/MaxHalford/chime) to use `@chime_on_end`) 67 | 68 | ## License 69 | 70 | The project has a [MIT Licence](https://github.com/FBruzzesi/deczoo/blob/main/LICENSE) 71 | -------------------------------------------------------------------------------- /tests/decorators_test/shape_tracker_test.py: -------------------------------------------------------------------------------- 1 | from contextlib import nullcontext as does_not_raise 2 | 3 | import numpy as np 4 | import pytest 5 | 6 | from deczoo import shape_tracker 7 | from deczoo._utils import EmptyShapeError 8 | 9 | 10 | @pytest.mark.parametrize( 11 | "arg_name, value, context", 12 | [ 13 | ("shape_in", True, does_not_raise()), 14 | ("shape_out", True, does_not_raise()), 15 | ("shape_delta", True, does_not_raise()), 16 | ("raise_if_empty", True, does_not_raise()), 17 | ("shape_in", False, does_not_raise()), 18 | ("shape_out", False, does_not_raise()), 19 | ("shape_delta", False, does_not_raise()), 20 | ("raise_if_empty", False, does_not_raise()), 21 | ("arg_to_track", 0, does_not_raise()), 22 | ("arg_to_track", "a", does_not_raise()), 23 | ("logging_fn", print, does_not_raise()), 24 | ("shape_in", 42, pytest.raises(TypeError)), 25 | ("shape_out", 1.0, pytest.raises(TypeError)), 26 | ("shape_delta", "a", pytest.raises(TypeError)), 27 | ("raise_if_empty", {1: 2}, pytest.raises(TypeError)), 28 | ("arg_to_track", (1, 2), pytest.raises(TypeError)), 29 | ("logging_fn", [1, 2], pytest.raises(TypeError)), 30 | ], 31 | ) 32 | def test_params(base_add, arg_name, value, context): 33 | """Tests that shape_tracker raises TypeError when invalid params are passed""" 34 | with context: 35 | shape_tracker(base_add, **{arg_name: value}) 36 | 37 | 38 | # Since we tested for param errors, now we can just test for every boolean turned on 39 | @pytest.mark.parametrize( 40 | "arg_to_track, n, context", 41 | [ 42 | (0, 3, does_not_raise()), 43 | ("a", 3, does_not_raise()), 44 | (0, 0, pytest.raises(EmptyShapeError)), 45 | ("a", 0, pytest.raises(EmptyShapeError)), 46 | (1, 3, pytest.raises(AttributeError)), 47 | ("b", 3, pytest.raises(KeyError)), 48 | ], 49 | ) 50 | def test_shape_in(capsys, arg_to_track, n, context): 51 | """Tests that shape_tracker tracks input shape""" 52 | a_shape = (1, 2) 53 | 54 | @shape_tracker( 55 | shape_in=True, 56 | shape_out=True, 57 | shape_delta=True, 58 | raise_if_empty=True, 59 | arg_to_track=arg_to_track, 60 | logging_fn=print, 61 | ) 62 | def n_vstack(a: np.ndarray, n: int) -> np.ndarray: 63 | if n > 0: 64 | return np.vstack(n * [a]) 65 | else: 66 | return np.empty(0) 67 | 68 | with context: 69 | n_vstack(np.ones(a_shape), n) 70 | sys_out = capsys.readouterr().out 71 | 72 | assert "Input: `a` has shape" in sys_out 73 | assert "Output: result has shape " in sys_out 74 | assert "Shape delta: " in sys_out 75 | -------------------------------------------------------------------------------- /docs/decorators/intro.md: -------------------------------------------------------------------------------- 1 | # Introduction 2 | 3 | ## What is a decorator? 4 | 5 | In Python, a decorator is a way to modify or enhance the behavior of a function or a class without actually modifying the source code of the function or class. 6 | 7 | Defining this _additional behaviour_ in a decorator, instead of within the original function itself, allows us to re-use it every time we need it, without a tedious _copy-pasting_. 8 | 9 | Decorators are implemented as functions (or classes) that take a function or a class as input and return a new function or class that has some additional functionality. 10 | 11 | Here is a simple example in order to illustrate the concept: 12 | 13 | ```python 14 | from typing import Callable 15 | 16 | def my_decorator(func: Callable) -> Callable: 17 | """ 18 | - `func` is the function taken as input to our decorator 19 | - `func` behaviour will be modified/enhanced 20 | - `wrapper` is the function that the decorator returns 21 | """ 22 | 23 | def wrapper(*args, **kwargs): 24 | 25 | print(f"Starting to run {func.__name__}") 26 | res = func(*args, **kwargs) 27 | print(f"{func.__name__} finished running!") 28 | 29 | return res 30 | 31 | return wrapper 32 | ``` 33 | 34 | Here, `my_decorator` is a function that takes a function as input, denoted as `func`, and returns a new function called `wrapper` defined within the inner scope. In this case, `wrapper` contains some additional functionalities (in this case, the `print` statements) that are executed before and after the original function call. 35 | 36 | `wrapper` function takes any number of positional and keyword arguments (`*args` and `**kwargs`), calls the original function `func` with those arguments, and returns the original result of the function call `res`. 37 | 38 | ## Decorator syntax 39 | 40 | We just saw what is a decorator and how to code one, but how can we use it? 41 | 42 | There are two equivalent options available: 43 | 44 | - _Functional assignment_, we assign the output of the decorator to a new function: 45 | 46 | ```python 47 | def my_func(): 48 | print("Hello world!") 49 | return 42 50 | 51 | my_func = my_decorator(my_func) 52 | ``` 53 | 54 | - The (more pythonic) _`@`-syntax_ (also called _decorator syntax_ or _decorator notation_), this is a shorthand way of applying a decorator to a function or a class, without having to explicitly call the decorator function and passing the function or class as an argument: 55 | 56 | ```python 57 | @my_decorator 58 | def my_func(): 59 | print("Hello world!") 60 | return 42 61 | ``` 62 | 63 | In both cases, calling the function results in the following: 64 | 65 | ```python 66 | result = my_func() 67 | # Starting to run my_func 68 | # Hello world! 69 | # my_func finished running! 70 | 71 | result 72 | # 42 73 | ``` 74 | -------------------------------------------------------------------------------- /.github/workflows/deploy-docs.yaml: -------------------------------------------------------------------------------- 1 | name: Deploy Documentation 2 | 3 | on: 4 | workflow_dispatch: 5 | 6 | permissions: 7 | contents: write 8 | 9 | jobs: 10 | deploy: 11 | runs-on: ubuntu-latest 12 | steps: 13 | - name: Checkout source code 14 | uses: actions/checkout@v5 15 | - name: Configure Git Credentials 16 | run: | 17 | git config user.name github-actions[bot] 18 | git config user.email 41898282+github-actions[bot]@users.noreply.github.com 19 | - name: Set up Python 20 | uses: actions/setup-python@v6 21 | with: 22 | python-version: "3.10" 23 | - run: echo "cache_id=$(date --utc '+%V')" >> $GITHUB_ENV 24 | - uses: actions/cache@v4 25 | with: 26 | key: mkdocs-material-${{ env.cache_id }} 27 | path: .cache 28 | restore-keys: | 29 | mkdocs-material- 30 | - name: Install dependencies and deploy 31 | run: | 32 | python -m pip install --upgrade pip --no-cache-dir 33 | python -m pip install ."[docs]" --no-cache-dir 34 | mkdocs gh-deploy --force 35 | 36 | 37 | # name: Deploy Documentation 38 | 39 | # on: 40 | # workflow_dispatch: 41 | # push: 42 | # branches: 43 | # - main 44 | 45 | # jobs: 46 | 47 | # deploy-docs: 48 | # runs-on: ubuntu-latest 49 | 50 | # steps: 51 | 52 | # - name: Checkout gh-pages 53 | # uses: actions/checkout@v5 54 | # with: 55 | # ref: gh-pages 56 | # fetch-depth: 0 # Fetch all history for all branches and tags. 57 | 58 | # - name: Fetch latest version 59 | # run : | 60 | # git fetch origin gh-pages 61 | # LATEST_VERSION=$(git show gh-pages:versions.json | python -c "import sys, json; print(max([d['version'] for d in json.load(sys.stdin) if d['version']!='latest']))") 62 | 63 | # - name: Checkout current branch 64 | # uses: actions/checkout@v5 65 | # with: 66 | # fetch-depth: 0 # Fetch all history for all branches and tags. 67 | 68 | # - name: Markup Link Checker (mlc) 69 | # uses: becheran/mlc@v0.16.3 70 | # with: 71 | # args: docs 72 | 73 | # - name: Set up Python 74 | # uses: actions/setup-python@v6 75 | # with: 76 | # python-version: "3.10" 77 | 78 | # - name: Install dependencies 79 | # run: | 80 | # python -m pip install --upgrade pip 81 | # python -m pip install ."[docs]" 82 | 83 | # - name: Extract versions and deploy if changed 84 | # run: | 85 | # CURRENT_VERSION=$(python -c "import iso_week_date; print(iso_week_date.__version__)") 86 | # if [ "$CURRENT_VERSION" != "$LATEST_VERSION" ]; then 87 | # mike deploy $CURRENT_VERSION latest 88 | # mike set-alias latest $CURRENT_VERSION 89 | # git config user.name 'GitHub Actions Bot' 90 | # git config user.email 'github-actions-bot@users.noreply.github.com' 91 | # mike push 92 | # fi 93 | -------------------------------------------------------------------------------- /tests/decorators_test/multi_shape_tracker_test.py: -------------------------------------------------------------------------------- 1 | from contextlib import nullcontext as does_not_raise 2 | from typing import Tuple 3 | 4 | import numpy as np 5 | import pytest 6 | 7 | from deczoo import multi_shape_tracker 8 | from deczoo._utils import EmptyShapeError 9 | 10 | a = np.ones((1, 2)) 11 | b = np.ones((1, 2)) 12 | 13 | 14 | def add_multi(a: np.ndarray, b: np.ndarray) -> Tuple[np.ndarray, np.ndarray]: 15 | """Function to test multi_shape_tracker""" 16 | return a + b, a * b 17 | 18 | 19 | @pytest.mark.parametrize("arg_name", ["shapes_in", "shapes_out", "raise_if_empty", "logging_fn"]) 20 | @pytest.mark.parametrize("value", [1.1, {1: 2}, ("a", 1), (0, "b"), "a"]) 21 | def test_invalid_args(arg_name, value): 22 | """Tests that multi_shape_tracker raises a TyperError any invalid args are passed""" 23 | if arg_name == "shapes_in" and value == "a": 24 | pytest.skip("Not of interest") 25 | 26 | with pytest.raises(TypeError): 27 | _ = multi_shape_tracker(add_multi, **{arg_name: value})(a, b) 28 | 29 | 30 | @pytest.mark.parametrize("value", ["a", "b", ("a", "b"), 0, 1, (0, 1), None]) 31 | def test_valid_shapes_in(capsys, value): 32 | """Tests that multi_shape_tracker doesn't raise an error for valid shapes_in param""" 33 | 34 | with does_not_raise(): 35 | _ = multi_shape_tracker(add_multi, shapes_in=value)(a, b) 36 | 37 | if value is not None: 38 | sys_out = capsys.readouterr().out 39 | assert "Input shapes: " in sys_out 40 | 41 | 42 | @pytest.mark.parametrize("value", ["all", 0, 1, (0, 1), None]) 43 | def test_valid_shapes_out(capsys, value): 44 | """Tests that multi_shape_tracker doesn't raise an error for valid shapes_out param""" 45 | 46 | with does_not_raise(): 47 | _ = multi_shape_tracker(add_multi, shapes_out=value)(a, b) 48 | 49 | if value is not None: 50 | sys_out = capsys.readouterr().out 51 | assert "Output shapes: " in sys_out 52 | 53 | 54 | @pytest.mark.parametrize("value", ["all", "any", None]) 55 | @pytest.mark.parametrize( 56 | "shapes, context", 57 | [((1, 2), does_not_raise()), ((0,), pytest.raises(EmptyShapeError))], 58 | ) 59 | def test_valid_raise_if_empty(value, shapes, context): 60 | """Tests that multi_shape_tracker doesn't raise an error for valid raise_if_empty""" 61 | 62 | a = np.ones(shapes) 63 | b = np.ones(shapes) 64 | 65 | # forcing to switch context if value is None 66 | context = context if value is not None else does_not_raise() 67 | 68 | with context: 69 | _ = multi_shape_tracker(add_multi, raise_if_empty=value)(a, b) 70 | 71 | 72 | def test_valid_logging_fn(): 73 | """Tests that multi_shape_tracker doesn't raise an error for valid logging_fn param""" 74 | 75 | with does_not_raise(): 76 | _ = multi_shape_tracker(add_multi, logging_fn=print)(a, b) 77 | 78 | 79 | # multi_shape_tracker( 80 | # func: Optional[Callable[[HasShape, Sequence[Any]], Tuple[HasShape, ...]]] = None, 81 | # shapes_in: Optional[Union[str, int, Sequence[str], Sequence[int]]] = None, 82 | # shapes_out: Optional[Union[int, Sequence[int], Literal["all"]]] = "all", 83 | # raise_if_empty: Optional[Literal["any", "all"]] = "any", 84 | # logging_fn: Callable = LOGGING_FN, 85 | # ) -> Callable: 86 | -------------------------------------------------------------------------------- /mkdocs.yml: -------------------------------------------------------------------------------- 1 | # Project information 2 | site_name: Deczoo 3 | site_url: https://fbruzzesi.github.io/deczoo/ 4 | site_author: Francesco Bruzzesi 5 | site_description: Deczoo - A zoo for decorators 6 | 7 | # Repository information 8 | repo_name: FBruzzesi/deczoo 9 | repo_url: https://github.com/fbruzzesi/deczoo 10 | edit_uri: edit/main/docs/ 11 | 12 | # Configuration 13 | watch: 14 | - deczoo 15 | use_directory_urls: true 16 | theme: 17 | name: material 18 | font: 19 | text: Ubuntu 20 | code: Ubuntu Mono 21 | highlightjs: true 22 | hljs_languages: 23 | - bash 24 | - python 25 | palette: 26 | - media: '(prefers-color-scheme: light)' 27 | scheme: default 28 | primary: teal 29 | accent: deep-orange 30 | toggle: 31 | icon: material/lightbulb 32 | name: Switch to light mode 33 | - media: '(prefers-color-scheme: dark)' 34 | scheme: slate 35 | primary: teal 36 | accent: deep-orange 37 | toggle: 38 | icon: material/lightbulb-outline 39 | name: Switch to dark mode 40 | features: 41 | - navigation.tabs 42 | - navigation.tabs.sticky 43 | - navigation.sections 44 | - navigation.expand 45 | - navigation.path 46 | - navigation.indexes 47 | - navigation.footer 48 | - navigation.top 49 | - navigation.tracking 50 | - content.action.edit 51 | - content.action.view 52 | - content.code.annotate 53 | - content.code.copy 54 | - content.tooltips 55 | - content.tabs.link 56 | - search.suggest 57 | - search.highlight 58 | - search.share 59 | - toc.follow 60 | - toc.integrate 61 | logo: img/deczoo-logo.png 62 | favicon: img/deczoo-logo.png 63 | 64 | # Plugins 65 | plugins: 66 | - mkdocstrings 67 | - search: 68 | separator: '[\s\-,:!=\[\]()"`/]+|\.(?!\d)|&[lg]t;|(?!\b)(?=[A-Z][a-z])' 69 | 70 | # Customization 71 | extra: 72 | social: 73 | - icon: fontawesome/brands/github 74 | link: https://github.com/fbruzzesi 75 | - icon: fontawesome/brands/linkedin 76 | link: https://www.linkedin.com/in/francesco-bruzzesi/ 77 | - icon: fontawesome/brands/python 78 | link: https://pypi.org/project/deczoo/ 79 | 80 | # Extensions 81 | markdown_extensions: 82 | - abbr 83 | - admonition 84 | - attr_list 85 | - codehilite 86 | - def_list 87 | - footnotes 88 | - md_in_html 89 | - toc: 90 | permalink: true 91 | - pymdownx.arithmatex: 92 | generic: true 93 | - pymdownx.betterem: 94 | smart_enable: all 95 | - pymdownx.caret 96 | - pymdownx.details 97 | - pymdownx.highlight: 98 | anchor_linenums: true 99 | - pymdownx.inlinehilite 100 | - pymdownx.keys 101 | - pymdownx.magiclink: 102 | repo_url_shorthand: true 103 | user: squidfunk 104 | repo: mkdocs-material 105 | - pymdownx.mark 106 | - pymdownx.smartsymbols 107 | - pymdownx.superfences 108 | - pymdownx.tabbed: 109 | alternate_style: true 110 | - pymdownx.tasklist: 111 | custom_checkbox: true 112 | - pymdownx.tilde 113 | 114 | nav: 115 | - Home: index.md 116 | - Decorators: 117 | - Introduction: decorators/intro.md 118 | - Advanced Usage: decorators/advanced.md 119 | - Quickstart: quickstart.md 120 | - API: 121 | - Decorators: api/decorators.md 122 | - Utils: api/utils.md 123 | - Contributing: contribute.md 124 | -------------------------------------------------------------------------------- /docs/contribute.md: -------------------------------------------------------------------------------- 1 | # Contributing 2 | 3 | ## Guidelines 4 | 5 | We welcome contributions to the library! If you have a bug fix or new feature that you would like to contribute, please follow the steps below: 6 | 7 | 1. Fork the repository on GitHub. 8 | 2. Clone the repository to your local machine. 9 | 3. Create a new branch for your bug fix or feature. 10 | 4. Make your changes and test them thoroughly, making sure that it passes all current tests. 11 | 5. Commit your changes and push the branch to your fork. 12 | 6. Open a pull request on the main repository. 13 | 14 | ## Code of Conduct 15 | 16 | All contributors are expected to follow the project's code of conduct, which is based on the Contributor Covenant. 17 | 18 | ### Reporting Bugs 19 | 20 | If you find a bug in the library, please report it by opening an [issue on GitHub](https://github.com/FBruzzesi/deczoo/issues). Be sure to include the version of the library you're using, as well as any error messages or tracebacks and a reproducible example. 21 | 22 | ### Requesting Features 23 | 24 | If you have a suggestion for a new feature, please open an [issue on GitHub](https://github.com/FBruzzesi/deczoo/issues). Be sure to explain the problem that you're trying to solve and how you think the feature would solve it. 25 | 26 | ### Submitting Pull Requests 27 | 28 | When submitting a pull request, please make sure that you've followed the steps above and that your code has been thoroughly tested. Also, be sure to include a brief summary of the changes you've made and a reference to any issues that your pull request resolves. 29 | 30 | ### Code formatting 31 | 32 | Compclasses uses [black](https://black.readthedocs.io/en/stable/index.html) and [isort](https://pycqa.github.io/isort/) with the following parameters for code formatting: 33 | 34 | ```bash 35 | isort --profile black -l 90 deczoo tests 36 | black --target-version py38 --line-length 90 deczoo tests 37 | ``` 38 | 39 | As part of the checks on pull requests, it is checked whether the code follows those standards. To ensure that the standard is met, it is recommended to install [pre-commit hooks](https://pre-commit.com/): 40 | 41 | ```bash 42 | python -m pip install pre-commit 43 | pre-commit install 44 | ``` 45 | 46 | ## Developing 47 | 48 | Let's suppose that you already did steps 1-3 from the above list, now you should install the library and its developing dependencies in editable way. 49 | 50 | First move into the repo folder: `cd deczoo`. 51 | 52 | Then: 53 | 54 | === "with make" 55 | 56 | ```bash 57 | make init-dev 58 | ``` 59 | 60 | === "without make" 61 | 62 | ```bash 63 | pip install -e ".[all]" --no-cache-dir 64 | pre-commit install 65 | ``` 66 | 67 | Now you are ready to proceed with all the changes you want to! 68 | 69 | ## Testing 70 | 71 | Once you are done with changes, you should: 72 | 73 | - add tests for the new features in the `/tests` folder 74 | - make sure that new features do not break existing codebase by running tests: 75 | 76 | === "with make" 77 | 78 | ```bash 79 | make test 80 | ``` 81 | 82 | === "without make" 83 | 84 | ```bash 85 | pytest tests -vv 86 | ``` 87 | 88 | ## Docs 89 | 90 | The documentation is generated using [mkdocs-material](https://squidfunk.github.io/mkdocs-material/), the API part uses [mkdocstrings](https://mkdocstrings.github.io/). 91 | 92 | If a breaking feature is developed, then we suggest to update documentation in the `/docs` folder as well, in order to describe how this can be used from a user perspective. 93 | -------------------------------------------------------------------------------- /deczoo/_utils.py: -------------------------------------------------------------------------------- 1 | import sys 2 | from functools import partial, wraps 3 | from typing import Callable, Protocol, Tuple, TypeVar, Union, runtime_checkable 4 | 5 | if sys.version_info >= (3, 10): 6 | from typing import ParamSpec, TypeAlias 7 | else: 8 | from typing_extensions import ParamSpec, TypeAlias 9 | 10 | if sys.version_info >= (3, 11): 11 | from typing import Self 12 | else: 13 | from typing_extensions import Self 14 | 15 | PS = ParamSpec("PS") 16 | R = TypeVar("R") 17 | F: TypeAlias = Callable[PS, R] 18 | 19 | # DPS = ParamSpec("DPS") 20 | # DecoratorType: TypeAlias = Callable[[F, DPS.args, DPS.kwargs], F] 21 | 22 | 23 | def check_parens(decorator: F) -> F: 24 | """Check whether or not a decorator function gets called with parens: 25 | 26 | - If called with parens, the decorator is called without the function as the first argument, but necessarely with 27 | decorator keyword arguments. 28 | - If called without parens, the decorator is called with the function as the first argument, and the decorator's 29 | default arguments. 30 | 31 | This function is used internally to endow every decorator of the above property. 32 | 33 | Arguments: 34 | decorator: decorator to wrap 35 | 36 | Returns: 37 | Wrapped decorator. 38 | 39 | Usage: 40 | ```python 41 | @check_parens 42 | def decorator(func, k1="default1", k2="default2"): 43 | # where the magic happens 44 | ... 45 | 46 | # `decorator` called without parens, hence default params. 47 | @decorator 48 | def func(*func_args, **func_kwargs): 49 | pass 50 | 51 | # `decorator` called with custom params, necessarely using parens. 52 | @decorator(*args, **kwargs) 53 | def func(*func_args, **func_kwargs): 54 | pass 55 | ``` 56 | """ 57 | 58 | @wraps(decorator) 59 | def wrapper(func: Union[F, None] = None, *args: PS.args, **kwargs: PS.kwargs) -> F: 60 | if func is None: 61 | return partial(decorator, *args, **kwargs) 62 | else: 63 | return decorator(func, *args, **kwargs) 64 | 65 | return wrapper 66 | 67 | 68 | def _get_free_memory() -> int: 69 | """Computes machine free memory via `/proc/meminfo` file (linux only). 70 | 71 | !!! warning 72 | This functionality is supported on unix-based systems only! 73 | """ 74 | 75 | with open("/proc/meminfo", "r") as mem: 76 | free_memory = 0 77 | for i in mem: 78 | sline = i.split() 79 | if str(sline[0]) in ("MemFree:", "Buffers:", "Cached:"): 80 | free_memory += int(sline[1]) 81 | return free_memory 82 | 83 | 84 | @runtime_checkable 85 | class SupportShape(Protocol): 86 | """Protocol for objects that have a `.shape()` attribute. In this context, a dataframe or array like object.""" 87 | 88 | @property 89 | def shape(self: Self) -> Tuple[int, ...]: 90 | ... 91 | 92 | 93 | class EmptyShapeError(Exception): 94 | """Exception raised when a dataframe/array-like object has an empty shape.""" 95 | 96 | ... 97 | 98 | 99 | LoggerType: TypeAlias = Callable[[str], None] 100 | LOGGING_FN: LoggerType 101 | 102 | try: 103 | from rich.console import Console 104 | from rich.theme import Theme 105 | 106 | custom_theme = Theme({"good": "bold green", "bad": "bold red"}) 107 | console = Console(theme=custom_theme) 108 | 109 | LOGGING_FN = console.log 110 | 111 | except ImportError: 112 | LOGGING_FN = print 113 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | 2 | 3 | ![](https://img.shields.io/github/license/FBruzzesi/deczoo) 4 | 5 | [![Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff) 6 | 7 | 8 | 9 | # Deczoo 10 | 11 | > A zoo for decorators 12 | 13 | There are many great decorators out there that we use everyday. Why don't collect few of them? 14 | 15 | I found myself implementing over and over in different projects. The hope is to gather them here and use this codebase. 16 | 17 | --- 18 | 19 | [Documentation](https://fbruzzesi.github.io/deczoo/) | [Source Code](https://github.com/fbruzzesi/deczoo) 20 | 21 | --- 22 | 23 | ## Alpha Notice 24 | 25 | This codebase is experimental and is working for my use cases. It is very probable that there are cases not covered and for which it breaks (badly). If you find them, please feel free to open an issue in the [issue page](https://github.com/FBruzzesi/deczoo/issues) of the repo. 26 | 27 | ## What is a decorator? 28 | 29 | In short a python decorator is a way to modify or enhance the behavior of a function or a class without actually modifying the source code of the function or class. 30 | 31 | Decorators are implemented as functions (or classes) that take a function or a class as input and return a new function or class that has some additional functionality. 32 | 33 | To have a more in-depth explanation you can check the [decorators docs section](https://fbruzzesi.github.io/deczoo/decorators/intro/). 34 | 35 | ## Installation 36 | 37 | **deczoo** is published as a Python package on [pypi](https://pypi.org/), and it can be installed with pip, or directly from source using git, or with a local clone: 38 | 39 | - **pip** (suggested): 40 | 41 | ```bash 42 | python -m pip install deczoo 43 | ``` 44 | 45 | - **pip + source/git**: 46 | 47 | ```bash 48 | python -m pip install git+https://github.com/FBruzzesi/deczoo.git 49 | ``` 50 | 51 | - **local clone**: 52 | 53 | ```bash 54 | git clone https://github.com/FBruzzesi/deczoo.git 55 | cd deczoo 56 | python -m pip install . 57 | ``` 58 | 59 | ### Dependencies 60 | 61 | As of now, the library has no additional required dependencies, however: 62 | 63 | - some functionalities works only on UNIX systems (`@memory_limit` and `@timeout`) 64 | - to use some decorators you may need to install additional dependencies (e.g. install [`chime`](https://github.com/MaxHalford/chime) to use `@chime_on_end`) 65 | 66 | ## Getting started 67 | 68 | The idea is kind of simple: each function in the library is a (function) decorator with a specific objective in mind. 69 | 70 | ```python title="Example: log decorator" 71 | from deczoo import log 72 | 73 | @log # equivalent to @log(log_time=True, log_args=True, log_error=True, logging_fn=print) 74 | def custom_add(a, b, *args): 75 | """Adds all arguments together""" 76 | return sum([a, b, *args]) 77 | 78 | _ = custom_add(1, 2, 3, 4) 79 | # custom_add args=(a=1, b=2, args=(3, 4)) time=0:00:00.000062 80 | 81 | _ = custom_add(1, "a", 2) 82 | # custom_add args=(a=1, b=a, args=(2,)) time=0:00:00.000064 Failed with error: unsupported 83 | # operand type(s) for +: 'int' and 'str' 84 | ``` 85 | 86 | ```python title="Example: shape_tracker decorator" 87 | from deczoo import shape_tracker 88 | 89 | @shape_tracker(shape_in=True, shape_out=True, shape_delta=True, raise_if_empty=True) 90 | def tracked_vstack(a: np.ndarray, b: np.ndarray) -> np.ndarray: 91 | return np.vstack([a, b]) 92 | 93 | _ = tracked_vstack(np.ones((1, 2)), np.ones((10, 2))) 94 | # Input: `a` has shape (1, 2) 95 | # Output: result has shape (11, 2) 96 | # Shape delta: (-10, 0) 97 | ``` 98 | 99 | ### Features 100 | 101 | The library implements the following decorators: 102 | 103 | - `call_counter`: tracks how many times a function has been called. 104 | - `catch`: wraps a function in a try-except block, returning a custom value, or raising a custom exception. 105 | - `check_args`: checks that function arguments satisfy its "rule". 106 | - `chime_on_end`: notify with chime sound on function end (success or error). 107 | - `log`: tracks function time taken, arguments and errors, such logs can be written to a file. 108 | - `timer`: tracks function time taken. 109 | - `memory_limit`: sets a memory limit while running the function. 110 | - `notify_on_end`: notifies when function finished running with a custom notifier. 111 | - `raise_if`: raises a custom exception if a condition is met. 112 | - `retry`: wraps a function with a "retry" block. 113 | - `shape_tracker`: tracks the shape of a dataframe/array-like object, in input and/or output. 114 | - `multi_shape_tracker`: tracks the shapes of input(s) and/or output(s) of a function. 115 | - `timeout`: sets a time limit for the function, terminates the process if it hasn't finished within such time limit. 116 | 117 | ### Examples 118 | 119 | Please refer to the [api page](https://fbruzzesi.github.io/deczoo/api/decorators/) to see a basic example for each decorator. 120 | 121 | ## Contributing 122 | 123 | Please read the [Contributing guidelines](https://fbruzzesi.github.io/deczoo/contribute/) in the documentation site. 124 | 125 | ## License 126 | 127 | The project has a [MIT Licence](https://github.com/FBruzzesi/deczoo/blob/main/LICENSE) 128 | -------------------------------------------------------------------------------- /docs/decorators/advanced.md: -------------------------------------------------------------------------------- 1 | # Advanced Usage 2 | 3 | ## Stacks 4 | 5 | We're not limited to a single decorator per function, we can _stack_ how many we want. 6 | 7 | Let's see how that works: 8 | 9 | ```python 10 | def decorator1(func): 11 | def wrapper(*args, **kwargs): 12 | 13 | print("Decorator 1") 14 | return func(*args, **kwargs) 15 | 16 | return wrapper 17 | 18 | def decorator2(func): 19 | def wrapper(*args, **kwargs): 20 | 21 | print("Decorator 2") 22 | return func(*args, **kwargs) 23 | 24 | return wrapper 25 | ``` 26 | 27 | Remark that the _order_ in which we stack decorators matter: 28 | 29 | === "\@d1 \@d2" 30 | 31 | ```python hl_lines="1 2 8 9" 32 | 33 | @decorator1 34 | @decorator2 35 | def func(): 36 | print("Hello world!") 37 | return 42 38 | 39 | func() 40 | # Decorator 1 41 | # Decorator 2 42 | # Hello world! 43 | ``` 44 | 45 | === "\@d2 \@d1" 46 | 47 | ```python hl_lines="1 2 8 9" 48 | 49 | @decorator2 50 | @decorator1 51 | def func(): 52 | print("Hello world!") 53 | return 42 54 | 55 | func() 56 | # Decorator 2 57 | # Decorator 1 58 | # Hello world! 59 | ``` 60 | 61 | ## Wraps 62 | 63 | [`functools.wraps`](https://docs.python.org/3/library/functools.html#functools.wraps) is a utility function in the Python standard library that is often used in decorators to preserve the original function's metadata (such as its name, docstring, and annotations) in the wrapper function. 64 | 65 | Here's an example of how `functools.wraps` can be used in a decorator, and how the metadata's differ: 66 | 67 | === "with \@wraps" 68 | 69 | ```python hl_lines="1 4 18 19" 70 | from functools import wraps 71 | 72 | def dec_with_wraps(func): 73 | @wraps(func) 74 | def wrapper(*args, **kwargs): 75 | # Do something before the function is called 76 | result = func(*args, **kwargs) 77 | # Do something after the function is called 78 | return result 79 | return wrapper 80 | 81 | @dec_with_wraps 82 | def my_func(): 83 | """This is my function""" 84 | return "Hello, world!" 85 | 86 | print(f"Name = '{my_func.__name__}'", f"Docs = '{my_func.__doc__}'", sep="\n") 87 | # Name = 'my_func' 88 | # Docs = 'This is my function' 89 | ``` 90 | 91 | === "without \@wraps" 92 | 93 | ```python hl_lines="15 16" 94 | def dec_no_wraps(func): 95 | def wrapper(*args, **kwargs): 96 | # Do something before the function is called 97 | result = func(*args, **kwargs) 98 | # Do something after the function is called 99 | return result 100 | return wrapper 101 | 102 | @dec_no_wraps 103 | def my_func(): 104 | """This is my function""" 105 | return "Hello, world!" 106 | 107 | print(f"Name = '{my_func.__name__}'", f"Docs = '{my_func.__doc__}'", sep="\n") 108 | # Name = 'wrapper' 109 | # Docs = 'None' 110 | ``` 111 | 112 | As you can see, the two decorators `dec_with_wraps` and `dec_no_wraps` are identical; the only difference between the two cases is the use of `@wraps` decorator in the former to preserve the metadata of the original function. 113 | 114 | When we print the `__name__` and `__doc__` attributes of the function in the two different scenarios, we obtain completely different results! In particular, in the first case the metadata of the decorated `my_func` are maintained, in the latter the metadata we obtain are those of the `wrapper` function inside the decorator. 115 | 116 | ## Decorators with arguments 117 | 118 | 119 | 120 | Sometimes we have more complexity to model and to achieve that we need to be able to pass arguments to our decorator. 121 | 122 | Let's assume that we want to run a function twice, or 3-times, or 4-times and so on. 123 | 124 | Instead of writing different decorators that run the input function N times, we can go one level deeper, and define a function that takes the decorator arguments and returns the actual decorator function. 125 | 126 | ```python title="repeat_n_times" 127 | from functools import wraps 128 | from typing import Callable 129 | 130 | def repeat_n_times(n: int) -> Callable: 131 | """Gets as input the arguments to be used in the actual decorator""" 132 | 133 | def decorator(func: Callable) -> Callable: 134 | """This is the actual decorator!""" 135 | 136 | @wraps(func) 137 | def wrapper(*args, **kwargs): 138 | """Returns a list with N func results""" 139 | return [func(*args, **kwargs) for _ in range(n)] 140 | 141 | return wrapper 142 | 143 | return decorator 144 | 145 | @repeat_n_times(n=2) 146 | def say_hello(name: str) -> str: 147 | return f"Hello {name}!" 148 | 149 | print(say_hello("Fra")) 150 | # ['Hello Fra!', 'Hello Fra!'] 151 | ``` 152 | 153 | 154 | 155 | Do you feel confused? If the answer is yes, it is because it is kinda confusing! 156 | 157 | A decorator with arguments is a function that takes arguments and returns another function that acts as the *actual* decorator. 158 | 159 | This decorator function takes a function as an argument and returns a new function that modifies the original function in some way. 160 | 161 | The key difference between a decorator with arguments and a regular decorator is that the decorator with arguments has an extra layer of nested functions. The outer function takes the arguments and returns the actual decorator function, while the inner function takes the original function as an argument and returns the modified function. 162 | 163 | Remark that even if we define `repeat_n_times` to have a default value for `n`, when we decorate a function we need to _call_ the decorator, since that returns the actual decorator that we want, namely we need to: 164 | 165 | ```python 166 | def repeat_n_times(n: int = 3) -> Callable: 167 | ... 168 | 169 | @repeat_n_times() 170 | def say_hello(name: str) -> str: 171 | return f"Hello {name}!" 172 | 173 | @repeat_n_times 174 | def say_goodbye(name: str) -> str: 175 | return f"Goodbye {name}!" 176 | 177 | 178 | print(say_hello("Fra")) 179 | # ['Hello Fra!', 'Hello Fra!', 'Hello Fra!'] 180 | 181 | print(say_goodbye("Fra")) 182 | # .decorator..wrapper> 183 | ``` 184 | 185 | Which is not really what we want for the `say_goodbye` function! 186 | 187 | Can we do it differently??? Sure we can! And that's how all decorators in **deczoo** are implemented. 188 | 189 | ## Decorators with arguments, and a trick! 190 | 191 | In the [introduction](intro.md) we saw how a decorator is defined, let's stuck to such implementation but let's see how to add additional parameters and control flow without the need to have more level of indentation. 192 | 193 | Here is a different implementation of `repeat_n_times`, this time without a triple level of indentation: 194 | 195 | ```python title="repeat_n_times definition" 196 | from functools import wraps, partial 197 | from typing import Callable 198 | 199 | def repeat_n_times(func: Callable = None, n: int = 2) -> Callable: 200 | 201 | @wraps(func) 202 | def wrapper(*args, **kwargs): 203 | 204 | results = [func(*args, **kwargs) for _ in range(n)] 205 | 206 | return results 207 | 208 | if func is None: 209 | return partial(repeat_n_times, n=n) 210 | else: 211 | return wrapper 212 | ``` 213 | 214 | Let's see what happens here: 215 | 216 | - `repeat_n_times` takes as input the function to decorate (`func`) as first argument, and any additional input right after. 217 | - To use this trick, every additional argument must have a default value (which can be `None`). 218 | - Within the decorator we implement a `wrapper` as usual, where we use any additional decorator argument (`n` in this example). 219 | - Since `wrapper` is not run until execution time, we can then check what is the value of `func`: 220 | - if it is `None`, then it means that only additional arguments have been provided to the decorator, and therefore we return a [partial](https://docs.python.org/3/library/functools.html#functools.partial) decorator with the given arguments that will decorate our function. 221 | - Otherwise, the function is provided and we return the `wrapper`. 222 | 223 | This "trick" allows us to use the decorator with parens, providing custom arguments, or without parens, using defaults, i.e. 224 | 225 | ```python title="repeat_n_times example" 226 | 227 | @repeat_n_times(n=3) # uses custom argument value 228 | def say_hello(name: str) -> str: 229 | return f"Hello {name}!" 230 | 231 | @repeat_n_times # uses default argument value 232 | def say_goodbye(name: str) -> str: 233 | return f"Goodbye {name}!" 234 | 235 | print(say_hello("Fra")) 236 | # ["Hello Fra!", "Hello Fra!", "Hello Fra!"] 237 | 238 | print(say_goodbye("Fra")) 239 | # ['Goodbye Fra!', 'Goodbye Fra!'] 240 | ``` 241 | 242 | Neat! This was possible to achieve using the control flow and `partial` block at the end of the decorator. 243 | 244 | Since in [deczoo](../index.md) every decorator is implemented using this strategy, we wrote a sort of "meta-decorator", called [check_parens](../api/utils.md#check_parens) that adds such block to every decorator! 245 | 246 | ```python 247 | ... 248 | if func is None: 249 | return partial(repeat_n_times, n=n) 250 | else: 251 | return wrapper 252 | ``` 253 | -------------------------------------------------------------------------------- /docs/img/interrogate-shield.svg: -------------------------------------------------------------------------------- 1 | 2 | interrogate: 100% 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | interrogate 14 | interrogate 15 | 100% 16 | 100% 17 | 18 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | 26 | 27 | 28 | 29 | 30 | 31 | 32 | 33 | 34 | 35 | 36 | 37 | 38 | 39 | 40 | 41 | 42 | 43 | 44 | 45 | 46 | 47 | 48 | 49 | 50 | 51 | 52 | 53 | 54 | 55 | 56 | 57 | 58 | 59 | -------------------------------------------------------------------------------- /deczoo/decorators.py: -------------------------------------------------------------------------------- 1 | import inspect 2 | import resource 3 | import signal 4 | import sys 5 | import time 6 | from enum import Enum 7 | from functools import partial, wraps 8 | from itertools import zip_longest 9 | from pathlib import Path 10 | from typing import Any, Callable, List, Literal, Sequence, Tuple, Type, TypeVar, Union 11 | 12 | from deczoo._base_notifier import BaseNotifier 13 | from deczoo._utils import ( 14 | LOGGING_FN, 15 | EmptyShapeError, 16 | SupportShape, 17 | _get_free_memory, 18 | check_parens, 19 | ) 20 | 21 | if sys.version_info >= (3, 10): 22 | from typing import ParamSpec 23 | else: 24 | from typing_extensions import ParamSpec 25 | 26 | PS = ParamSpec("PS") 27 | R = TypeVar("R") 28 | RE = TypeVar("RE") 29 | 30 | 31 | @check_parens 32 | def call_counter( 33 | func: Union[Callable[PS, R], None] = None, 34 | seed: int = 0, 35 | log_counter: bool = True, 36 | logging_fn: Callable[[str], None] = LOGGING_FN, 37 | ) -> Callable[PS, R]: 38 | """Counts how many times a function has been called by setting and tracking a `_calls` attribute to the decorated 39 | function. 40 | 41 | `_calls` is set from a given `seed` value, and incremented by 1 each time the function is called. 42 | 43 | Arguments: 44 | func: Function to decorate 45 | seed: Counter start 46 | log_counter: Whether or not to log `_calls` value each time the function is called 47 | logging_fn: Log function (e.g. print, logger.info, rich console.print) 48 | 49 | Raises: 50 | TypeError: If `seed` is not an int, `log_counter` is not a bool, or `logging_fn` is not a callable when 51 | `log_counter` is True. 52 | 53 | Returns: 54 | Decorated function 55 | 56 | Usage: 57 | ```python 58 | from deczoo import call_counter 59 | 60 | @call_counter(seed=0, log_counter=False) 61 | def add(a, b): 62 | return a+b 63 | 64 | for _ in range(3): 65 | add(1,2) 66 | 67 | add._calls 68 | 3 69 | ``` 70 | """ 71 | 72 | if not isinstance(seed, int): 73 | raise TypeError("`seed` argument must be an int") 74 | 75 | if not isinstance(log_counter, bool): 76 | raise TypeError("`log_counter` argument must be a bool") 77 | 78 | if (log_counter is True) and (not callable(logging_fn)): 79 | raise TypeError("`logging_fn` argument must be a callable") 80 | 81 | @wraps(func) # type: ignore 82 | def wrapper(*args: PS.args, **kwargs: PS.kwargs) -> R: 83 | wrapper._calls += 1 # type: ignore 84 | 85 | if log_counter: 86 | logging_fn(f"{func.__name__} called {wrapper._calls} times") # type: ignore 87 | 88 | return func(*args, **kwargs) # type: ignore 89 | 90 | # set counter dynamically 91 | wrapper._calls = seed # type: ignore 92 | 93 | return wrapper 94 | 95 | 96 | @check_parens 97 | def catch( 98 | func: Union[Callable[PS, R], None] = None, 99 | return_on_exception: Union[RE, None] = None, 100 | raise_on_exception: Union[Type[Exception], None] = None, 101 | logging_fn: Callable[[str], None] = LOGGING_FN, 102 | ) -> Callable[PS, Union[R, RE]]: 103 | """Wraps a function in a try-except block, potentially prevent exception to be raised by returning a given value or 104 | raises custom exception. 105 | 106 | Remark that if both `return_on_exception` and `raise_on_exception` are provided, `return_on_exception` will be used. 107 | 108 | Arguments: 109 | func: Function to decorate 110 | return_on_exception: Value to return on exception 111 | raise_on_exception: Error to raise on exception 112 | logging_fn: Log function (e.g. print, logger.info, rich console.print) 113 | 114 | Returns: 115 | Decorated function 116 | 117 | Raises: 118 | TypeError: If `logging_fn` is not a callable 119 | 120 | Usage: 121 | ```python 122 | from deczoo import catch 123 | 124 | @catch(return_on_exception=-999) 125 | def add(a, b): 126 | return a+b 127 | 128 | add(1, 2) 129 | 3 130 | 131 | add(1, "a") 132 | -999 133 | ``` 134 | """ 135 | 136 | if not callable(logging_fn): 137 | raise TypeError("`logging_fn` argument must be a callable") 138 | 139 | @wraps(func) # type: ignore 140 | def wrapper(*args: PS.args, **kwargs: PS.kwargs) -> Union[R, RE]: 141 | try: 142 | return func(*args, **kwargs) # type: ignore 143 | 144 | except Exception as e: 145 | if return_on_exception is not None: 146 | logging_fn(f"Failed with error {e}, returning {return_on_exception}") 147 | return return_on_exception 148 | 149 | elif raise_on_exception is not None: 150 | logging_fn(f"Failed with error {e}") 151 | raise raise_on_exception 152 | 153 | else: 154 | logging_fn(f"Failed with error {e}") 155 | raise e 156 | 157 | return wrapper 158 | 159 | 160 | @check_parens 161 | def check_args(func: Union[Callable[PS, R], None] = None, **rules: Callable[[Any], bool]) -> Callable[PS, R]: 162 | """Checks that function arguments satisfy given rules, if not a `ValueError` is raised. 163 | 164 | Each `rule` should be a keyword argument with the name of the argument to check, and the value should be a 165 | function/callable that takes the argument value and returns a boolean. 166 | 167 | Arguments: 168 | func: Function to decorate 169 | rules: Rules to be satisfied, each rule is a callable that takes the argument value and returns a boolean 170 | 171 | Returns: 172 | Decorated function 173 | 174 | Raises: 175 | ValueError: If any rule is not a callable 176 | ValueError: If any decorated function argument doesn't satisfy its rule 177 | 178 | Usage: 179 | ```python 180 | from deczoo import check_args 181 | 182 | @check_args(a=lambda t: t>0) 183 | def add(a, b): 184 | return a+b 185 | 186 | add(1,2) 187 | 3 188 | 189 | add(-2, 2) 190 | # ValueError: Argument `a` doesn't satisfy its rule 191 | ``` 192 | """ 193 | if not all(callable(rule) for rule in rules.values()): 194 | raise ValueError("All rules must be callable") 195 | 196 | @wraps(func) # type: ignore 197 | def wrapper(*args: PS.args, **kwargs: PS.kwargs) -> R: 198 | func_args = inspect.signature(func).bind(*args, **kwargs).arguments # type: ignore 199 | 200 | for k, v in func_args.items(): 201 | rule = rules.get(k) 202 | 203 | if rule is not None: 204 | if not rule(v): 205 | raise ValueError(f"Argument `{k}` doesn't satisfy its rule") 206 | 207 | return func(*args, **kwargs) # type: ignore 208 | 209 | return wrapper 210 | 211 | 212 | @check_parens 213 | def chime_on_end(func: Union[Callable[PS, R], None] = None, theme: str = "mario") -> Callable[PS, R]: 214 | """Notify with [chime](https://github.com/MaxHalford/chime) sound when function ends successfully or fails. 215 | 216 | Arguments: 217 | func: Function to decorate 218 | theme: Chime theme to use 219 | 220 | Returns: 221 | Decorated function 222 | 223 | Usage: 224 | ```python 225 | from deczoo import chime_on_end 226 | 227 | @chime_on_end 228 | def add(a, b): return a+b 229 | 230 | _ = add(1, 2) 231 | # you should hear a sound now! 232 | ``` 233 | """ 234 | import chime 235 | 236 | chime.theme(theme) 237 | 238 | @wraps(func) # type: ignore 239 | def wrapper(*args: PS.args, **kwargs: PS.kwargs) -> R: 240 | try: 241 | res = func(*args, **kwargs) # type: ignore 242 | chime.success() 243 | return res 244 | 245 | except Exception as e: 246 | chime.error() 247 | raise e 248 | 249 | return wrapper 250 | 251 | 252 | @check_parens 253 | def log( 254 | func: Union[Callable[PS, R], None] = None, 255 | log_time: bool = True, 256 | log_args: bool = True, 257 | log_error: bool = True, 258 | log_file: Union[Path, str, None] = None, 259 | logging_fn: Callable[[str], None] = LOGGING_FN, 260 | ) -> Callable[PS, R]: 261 | """Tracks function time taken, arguments and errors. If `log_file` is provided, logs are written to file. 262 | In any case, logs are passed to `logging_fn`. 263 | 264 | Arguments: 265 | func: Function to decorate 266 | log_time: Whether or not to track time taken 267 | log_args: Whether or not to track arguments 268 | log_error: Whether or not to track error 269 | log_file: Filepath where to write/save log string 270 | logging_fn: Log function (e.g. print, logger.info, rich console.print) 271 | 272 | Returns: 273 | Decorated function with logging capabilities 274 | 275 | Raises: 276 | TypeError: if `log_time`, `log_args` or `log_error` are not `bool` or `log_file` is not `None`, `str` or `Path` 277 | 278 | Usage: 279 | ```python 280 | from deczoo import log 281 | 282 | @log 283 | def add(a, b): return a+b 284 | 285 | _ = add(1, 2) 286 | # add args=(a=1, b=2) time=0:00:00.000111 287 | ``` 288 | """ 289 | 290 | if not all(isinstance(x, bool) for x in [log_time, log_args, log_error]): 291 | raise TypeError("`log_time`, `log_args` and `log_error` must be bool") 292 | 293 | if log_file is not None and not isinstance(log_file, (str, Path)): 294 | raise TypeError("`log_file` must be either None, str or Path") 295 | 296 | if not callable(logging_fn): 297 | raise TypeError("`logging_fn` must be callable") 298 | 299 | @wraps(func) # type: ignore 300 | def wrapper(*args: PS.args, **kwargs: PS.kwargs) -> R: 301 | tic = time.perf_counter() 302 | 303 | optional_strings: List[Union[str, None]] 304 | if log_args: 305 | func_args = inspect.signature(func).bind(*args, **kwargs).arguments # type: ignore 306 | func_args_str = ", ".join(f"{k}={v}" for k, v in func_args.items()) 307 | 308 | optional_strings = [f"args=({func_args_str})"] 309 | 310 | else: 311 | optional_strings = [] 312 | 313 | try: 314 | res = func(*args, **kwargs) # type: ignore 315 | toc = time.perf_counter() 316 | optional_strings += [ 317 | f"time={toc - tic}" if log_time else None, 318 | ] 319 | 320 | return res 321 | 322 | except Exception as e: 323 | toc = time.perf_counter() 324 | optional_strings += [ 325 | f"time={toc - tic}" if log_time else None, 326 | "Failed" + (f" with error: {e}" if log_error else ""), 327 | ] 328 | raise e 329 | 330 | finally: 331 | log_string = f"{func.__name__} {' '.join([s for s in optional_strings if s])}" # type: ignore 332 | logging_fn(log_string) 333 | 334 | if log_file is not None: 335 | with open(log_file, "a") as f: 336 | f.write(f"{tic} {log_string}\n") 337 | 338 | return wrapper 339 | 340 | 341 | timer = partial(log, log_time=True, log_args=False, log_error=False) 342 | 343 | 344 | @check_parens 345 | def memory_limit( 346 | func: Union[Callable[PS, R], None] = None, 347 | percentage: float = 0.99, 348 | logging_fn: Callable[[str], None] = LOGGING_FN, 349 | ) -> Callable[PS, R]: 350 | """Sets a memory limit while running the decorated function. 351 | 352 | !!! warning 353 | This decorator is supported on unix-based systems only! 354 | 355 | Arguments: 356 | func: Function to decorate 357 | percentage: Percentage of the currently available memory to use 358 | logging_fn: Log function (e.g. print, logger.info, rich console.print) 359 | 360 | Raises: 361 | TypeError: If `percentage` is not a `float` 362 | ValueError: If `percentage` is not between 0 and 1 363 | MemoryError: If memory limit is reached when decorated function is called 364 | 365 | Returns: 366 | Decorated function 367 | 368 | Usage: 369 | ```python 370 | from deczoo import memory_limit 371 | 372 | # Running on WSL2 with 12 Gb RAM 373 | 374 | @memory_limit(percentage=0.05) 375 | def limited(): 376 | for i in list(range(10 ** 8)): 377 | _ = 1 + 1 378 | return "done" 379 | 380 | def unlimited(): 381 | for i in list(range(10 ** 8)): 382 | _ = 1 + 1 383 | return "done" 384 | 385 | limited() 386 | # MemoryError: Reached memory limit 387 | 388 | unlimited() 389 | done 390 | ``` 391 | """ 392 | if not isinstance(percentage, float): 393 | raise TypeError("`percentage` should be a float") 394 | 395 | if not 0.0 <= percentage <= 1.0: 396 | raise ValueError("`percentage` should be between 0 and 1") 397 | 398 | if not callable(logging_fn): 399 | raise TypeError("`logging_fn` should be a callable") 400 | 401 | @wraps(func) # type: ignore 402 | def wrapper(*args: PS.args, **kwargs: PS.kwargs) -> R: 403 | _, hard = resource.getrlimit(resource.RLIMIT_AS) 404 | free_memory = _get_free_memory() * 1024 405 | 406 | logging_fn(f"Setting memory limit for {func.__name__} to {int(free_memory * percentage)}") # type: ignore 407 | 408 | resource.setrlimit(resource.RLIMIT_AS, (int(free_memory * percentage), hard)) 409 | 410 | try: 411 | return func(*args, **kwargs) # type: ignore 412 | 413 | except MemoryError: 414 | raise MemoryError("Reached memory limit") 415 | 416 | finally: 417 | resource.setrlimit(resource.RLIMIT_AS, (int(free_memory), hard)) 418 | 419 | return wrapper 420 | 421 | 422 | @check_parens 423 | def notify_on_end( 424 | func: Union[Callable[PS, R], None] = None, notifier: Union[BaseNotifier, None] = None 425 | ) -> Callable[PS, R]: 426 | """Notify when func has finished running using the notifier `notify` method. 427 | 428 | `notifier` object should inherit from BaseNotifier and implement any custom `.notify()` method. 429 | 430 | Arguments: 431 | func: Function to decorate 432 | notifier: Instance of a Notifier that implements `notify` method 433 | 434 | Returns: 435 | Decorated function 436 | 437 | Usage: 438 | ```python 439 | from deczoo import notify_on_end 440 | from deczoo._base_notifier import BaseNotifier 441 | 442 | class DummyNotifier(BaseNotifier): 443 | def notify(self): 444 | print("Function has finished") 445 | 446 | notifier = DummyNotifier() 447 | @notify_on_end(notifier=notifier) 448 | def add(a, b): return a + b 449 | 450 | _ = add(1, 2) 451 | # Function has finished 452 | ``` 453 | """ 454 | if not isinstance(notifier, BaseNotifier): 455 | raise TypeError("`notifier` should be an instance of a BaseNotifier") 456 | 457 | @wraps(func) # type: ignore 458 | def wrapper(*args: PS.args, **kwargs: PS.kwargs) -> R: 459 | try: 460 | return func(*args, **kwargs) # type: ignore 461 | except Exception as e: 462 | raise e 463 | finally: 464 | notifier.notify() 465 | 466 | return wrapper 467 | 468 | 469 | @check_parens 470 | def retry( 471 | func: Union[Callable[PS, R], None] = None, 472 | n_tries: int = 3, 473 | delay: float = 0.0, 474 | logging_fn: Callable[[str], None] = LOGGING_FN, 475 | ) -> Callable[PS, R]: 476 | """Wraps a function within a "retry" block. If the function fails, it will be retried `n_tries` times with a delay 477 | of `delay` seconds between each attempt. 478 | 479 | The function will be retried until it succeeds or the maximum number of attempts is reached. Either the first 480 | successful result will be returned or the last error will be raised. 481 | 482 | Arguments: 483 | func: Function to decorate 484 | n_tries: Max number of attempts to try 485 | delay: Time to wait before a retry 486 | logging_fn: Log function (e.g. print, logger.info, rich console.print) 487 | 488 | Raises: 489 | ValueError: If any of the following holds: 490 | - `n_tries` is not a positive integer 491 | - `delay` is not a positive number 492 | - `logging_fn` is not a callable 493 | 494 | Returns: 495 | Decorated function 496 | 497 | Usage: 498 | ```python 499 | from deczoo import retry 500 | 501 | @retry(n_tries=2, delay=1.) 502 | def add(a, b): return a+b 503 | 504 | _ = add(1, 2) 505 | # Attempt 1/2: Succeeded 506 | 507 | _ = add(1, "a") 508 | # Attempt 1/2: Failed with error: unsupported operand type(s) for +: 'int' and 'str' 509 | # Attempt 2/2: Failed with error: unsupported operand type(s) for +: 'int' and 'str' 510 | ``` 511 | """ 512 | if not isinstance(n_tries, int) or n_tries < 1: 513 | raise ValueError("`n_tries` should be a positive integer") 514 | 515 | if not isinstance(delay, (int, float)) or delay < 0: 516 | raise ValueError("`delay` should be a positive number") 517 | 518 | if not callable(logging_fn): 519 | raise TypeError("`logging_fn` should be a callable") 520 | 521 | @wraps(func) # type: ignore 522 | def wrapper(*args: PS.args, **kwargs: PS.kwargs) -> R: 523 | attempt = 0 524 | 525 | while attempt < n_tries: 526 | try: 527 | res = func(*args, **kwargs) # type: ignore 528 | logging_fn(f"Attempt {attempt+1}/{n_tries}: Succeeded") 529 | return res 530 | 531 | except Exception as e: 532 | logging_fn(f"Attempt {attempt+1}/{n_tries}: Failed with error: {e}") 533 | 534 | time.sleep(delay) 535 | attempt += 1 536 | if attempt == n_tries: 537 | raise e 538 | 539 | return wrapper 540 | 541 | 542 | @check_parens 543 | def shape_tracker( 544 | func: Union[Callable[[SupportShape, Sequence[Any]], SupportShape], None] = None, 545 | shape_in: bool = False, 546 | shape_out: bool = True, 547 | shape_delta: bool = False, 548 | raise_if_empty: bool = True, 549 | arg_to_track: Union[int, str] = 0, 550 | logging_fn: Callable[[str], None] = LOGGING_FN, 551 | ) -> Callable[[SupportShape, Sequence[Any]], SupportShape]: 552 | """Tracks the shape of a dataframe/array-like object. 553 | 554 | It's possible to track input and output shape(s), delta from input and output, and raise an error if resulting 555 | output is empty. 556 | 557 | This is particularly suitable to decorate functions that are used in a (pandas/polars/dask/...) pipe(line). 558 | 559 | The decorator will track the shape of the first argument of the function, unless `arg_to_track` is specified. 560 | 561 | `arg_to_track` can be: 562 | 563 | - a non-negative integer corresponding to the index of the argument to track 564 | - a string indicating the name of the argument to track. 565 | 566 | Parameters: 567 | func: Function to decorate 568 | shape_in: Track input shape 569 | shape_out: Track output shape 570 | shape_delta: Track shape delta between input and output 571 | raise_if_empty: Raise error if output is empty 572 | arg_to_track: Index or name of the argument to track, used only if `shape_in` is `True` 573 | logging_fn: Log function (e.g. print, logger.info, rich console.print) 574 | 575 | Returns: 576 | Decorated function 577 | 578 | Raises: 579 | TypeError: if any of the parameters is not of the correct type 580 | EmptyShapeError: if decorated function output is empty and `raise_if_empty` is `True` 581 | 582 | Usage: 583 | ```python 584 | import numpy as np 585 | from deczoo import shape_tracker 586 | 587 | @shape_tracker(shape_in=True, shape_out=True, shape_delta=True, raise_if_empty=True) 588 | def tracked_vstack(a: np.ndarray, b: np.ndarray) -> np.ndarray: 589 | return np.vstack([a, b]) 590 | 591 | _ = tracked_vstack(np.ones((1, 2)), np.ones((10, 2))) 592 | # Input: `a` has shape (1, 2) 593 | # Output: result has shape (11, 2) 594 | # Shape delta: (-10, 0) 595 | 596 | _ = tracked_vstack(np.ones((0, 2)), np.ones((0, 2))) 597 | # Input: `a` has shape (0, 2) 598 | # Output: result has shape (0, 2) 599 | # Shape delta: (0, 0) 600 | > EmptyShapeError: Result from tracked_vstack is empty 601 | ``` 602 | 603 | Now if the array to track is not the first argument, we can explicitly set 604 | `arg_to_track` to the value of 1 or "b". 605 | """ 606 | if not isinstance(shape_in, bool): 607 | raise TypeError("`shape_in` should be a boolean") 608 | 609 | if not isinstance(shape_out, bool): 610 | raise TypeError("`shape_out` should be a boolean") 611 | 612 | if not isinstance(shape_delta, bool): 613 | raise TypeError("`shape_delta` should be a boolean") 614 | 615 | if not isinstance(raise_if_empty, bool): 616 | raise TypeError("`raise_if_empty` should be a boolean") 617 | 618 | if (not isinstance(arg_to_track, (str, int))) or (isinstance(arg_to_track, int) and arg_to_track < 0): 619 | raise TypeError("`arg_to_track` should be a string or a positive integer") 620 | if not callable(logging_fn): 621 | raise TypeError("`logging_fn` should be a callable") 622 | 623 | @wraps(func) # type: ignore 624 | def wrapper(*args: PS.args, **kwargs: PS.kwargs) -> SupportShape: 625 | func_args = ( 626 | inspect.signature(func).bind(*args, **kwargs).arguments # type: ignore 627 | ) 628 | 629 | if isinstance(arg_to_track, int) and arg_to_track >= 0: 630 | _arg_name, _arg_value = tuple(func_args.items())[arg_to_track] 631 | elif isinstance(arg_to_track, str): 632 | _arg_name, _arg_value = arg_to_track, func_args[arg_to_track] 633 | else: 634 | raise ValueError("arg_to_track should be a string or a positive integer") 635 | 636 | if shape_in: 637 | logging_fn(f"Input: `{_arg_name}` has shape {_arg_value.shape}") 638 | 639 | res = func(*args, **kwargs) # type: ignore 640 | 641 | output_shape = res.shape 642 | 643 | if shape_out: 644 | logging_fn(f"Output: result has shape {output_shape}") 645 | 646 | if shape_delta: 647 | input_shape = _arg_value.shape 648 | delta = tuple(d1 - d2 for d1, d2 in zip_longest(input_shape, output_shape, fillvalue=0)) 649 | 650 | logging_fn(f"Shape delta: {delta}") 651 | 652 | if raise_if_empty and output_shape[0] == 0: 653 | raise EmptyShapeError(f"Result from {func.__name__} is empty") # type: ignore 654 | 655 | return res 656 | 657 | return wrapper 658 | 659 | 660 | @check_parens 661 | def multi_shape_tracker( 662 | func: Union[Callable[[SupportShape, Sequence[Any]], Tuple[SupportShape, ...]], None] = None, 663 | shapes_in: Union[str, int, Sequence[str], Sequence[int], None] = None, 664 | shapes_out: Union[int, Sequence[int], Literal["all"], None] = "all", 665 | raise_if_empty: Literal["any", "all", None] = "any", 666 | logging_fn: Callable[[str], None] = LOGGING_FN, 667 | ) -> Callable[[SupportShape, Sequence[Any]], Tuple[SupportShape, ...]]: 668 | """ 669 | Tracks the shape(s) of a dataframe/array-like objects both in input and output of 670 | a given function. 671 | 672 | This decorator differs from `shape_tracker` in that it can track the shape of 673 | multiple arguments and outputs. In particular it can track the shape of all the 674 | arguments and outputs of a function, and raise an error if _any_ of the tracked 675 | outputs is empty. 676 | 677 | Arguments: 678 | func: Function to decorate 679 | shapes_in: Sequence of argument positions OR argument names to track 680 | shapes_out: Sequence of output positions to track, "all" to track all, None to disable 681 | raise_if_empty: Raise error if tracked output results is/are empty (strategy: "any", "all", None) 682 | logging_fn: log function (e.g. print, logger.info, rich console.print) 683 | 684 | Returns: 685 | Decorated function 686 | 687 | Raises: 688 | TypeError: if any of the parameters is not of the correct type 689 | EmptyShapeError: if decorated function output is empty and `raise_if_empty` is `all` or `any` 690 | 691 | Usage: 692 | ```python 693 | import numpy as np 694 | from deczoo import multi_shape_tracker 695 | 696 | @multi_shape_tracker(shapes_in=(0, 1), shapes_out="all") 697 | def add_mult(a: np.ndarray, b: np.ndarray) -> Tuple[np.ndarray, ...]: 698 | return a + b, a * b, a@b.T 699 | 700 | a = b = np.ones((1, 2)) 701 | _ = add_mult(a, b) 702 | # Input shapes: a.shape=(1, 2) b.shape=(1, 2) 703 | # Output shapes: (1, 2) (1, 2) (1, 1) 704 | ``` 705 | `@multi_shape_tracker(shapes_in=(0, 1), shapes_out="all")` is equivalent to 706 | `@multi_shape_tracker(shapes_in=("a", "b"), shapes_out=(0, 1, 2))`. 707 | However we can choose to track a subset of inputs and outputs by using the 708 | `shapes_in` and `shapes_out` parameters. This is particularly useful when 709 | some inputs/outputs are not dataframe/array-like objects. 710 | """ 711 | 712 | if not callable(logging_fn): 713 | raise TypeError("`logging_fn` should be a callable") 714 | 715 | _arg_names: Union[str, Sequence[str]] 716 | _arg_values: Union[SupportShape, Sequence[SupportShape]] 717 | 718 | @wraps(func) # type: ignore 719 | def wrapper(*args: PS.args, **kwargs: PS.kwargs) -> Tuple[SupportShape, ...]: 720 | func_args = ( 721 | inspect.signature(func).bind(*args, **kwargs).arguments # type: ignore 722 | ) 723 | # parse shapes_in 724 | # case: str 725 | if isinstance(shapes_in, str): 726 | _arg_names, _arg_values = shapes_in, func_args[shapes_in] 727 | 728 | # case: int 729 | elif isinstance(shapes_in, int) and shapes_in >= 0: 730 | _arg_names, _arg_values = tuple(x for x in tuple(func_args.items())[shapes_in]) 731 | 732 | # case: sequence 733 | elif isinstance(shapes_in, Sequence): 734 | # case: sequence of str's 735 | if all(isinstance(x, str) for x in shapes_in): 736 | _arg_names, _arg_values = ( 737 | tuple(shapes_in), # type: ignore 738 | tuple(func_args[x] for x in shapes_in), # type: ignore 739 | ) 740 | 741 | # case: sequence of positive int's 742 | elif all(isinstance(x, int) and x >= 0 for x in shapes_in): 743 | _arg_names, _arg_values = zip( # type: ignore 744 | *(tuple(func_args.items())[x] for x in shapes_in) # type: ignore 745 | ) 746 | 747 | # case: sequence of something else! (raise error) 748 | else: 749 | raise TypeError("`shapes_in` values must all be str or positive int") 750 | 751 | # case: None 752 | elif shapes_in is None: 753 | pass 754 | 755 | # case: something else, not in Union[int, str, Sequence[int], Sequence[str], None] 756 | else: 757 | raise TypeError("`shapes_in` must be either a str, a positive int, a sequence of those or None") 758 | 759 | if shapes_in is not None: 760 | logging_fn("Input shapes: " + " ".join(f"{k}.shape={v.shape}" for k, v in zip(_arg_names, _arg_values))) 761 | 762 | # finally run the function! 763 | orig_res = func(*args, **kwargs) # type: ignore 764 | 765 | # Check if the function returns a single value or a tuple 766 | res = (orig_res,) if not isinstance(orig_res, Sequence) else orig_res 767 | 768 | # parse shapes_out 769 | # case: positive int 770 | if isinstance(shapes_out, int) and shapes_out >= 0: 771 | _res_shapes = (res[shapes_out].shape,) 772 | 773 | # case: sequence of positive int's 774 | elif isinstance(shapes_out, Sequence) and all(isinstance(x, int) and x >= 0 for x in shapes_out): 775 | _res_shapes = tuple(res[x].shape for x in shapes_out) # type: ignore 776 | 777 | # case: "all" 778 | elif shapes_out == "all": 779 | _res_shapes = tuple(x.shape for x in res) # type: ignore 780 | 781 | # case: None 782 | elif shapes_out is None: 783 | pass 784 | 785 | # case: something else, not in Union[int, Sequence[int], Literal["all"], None] 786 | else: 787 | raise TypeError("`shapes_out` must be positive int, sequence of positive int, 'all' or None") 788 | 789 | if shapes_out is not None: 790 | logging_fn("Output shapes: " + " ".join(f"{s}" for s in _res_shapes)) 791 | 792 | # parse raise_if_empty 793 | if (shapes_out is None) and (raise_if_empty is not None): # type: ignore 794 | _raise_if_empty = None 795 | 796 | logging_fn( 797 | "Overwriting `raise_if_empty` to None because `shapes_out` is None. " 798 | "Please specify `shapes_out` if you want to use `raise_if_empty`" 799 | ) 800 | else: 801 | _raise_if_empty = raise_if_empty 802 | 803 | # case: None 804 | if _raise_if_empty is None: 805 | pass 806 | # case: "any" 807 | elif _raise_if_empty == "any": 808 | if any(x[0] == 0 for x in _res_shapes): 809 | raise EmptyShapeError(f"At least one result from {func.__name__} is empty") # type: ignore 810 | # case: "all" 811 | elif _raise_if_empty == "all": 812 | if all(x[0] == 0 for x in _res_shapes): 813 | raise EmptyShapeError(f"All results from {func.__name__} are empty") # type: ignore 814 | 815 | # case: something else, not in Union[Literal["any"], Literal["all"], None] 816 | else: 817 | raise TypeError("raise_if_empty must be either 'any', 'all' or None") 818 | 819 | return orig_res # type: ignore 820 | 821 | return wrapper 822 | 823 | 824 | @check_parens 825 | def timeout( 826 | func: Union[Callable[PS, R], None] = None, 827 | time_limit: Union[int, None] = None, 828 | signal_handler: Union[Callable, None] = None, 829 | signum: Union[int, Enum] = signal.SIGALRM, 830 | ) -> Callable[PS, R]: 831 | """Sets a time limit to a function, terminates the process if it hasn't finished within such time limit. 832 | 833 | !!! warning 834 | This decorator uses the built-in [signal library](https://docs.python.org/3/library/signal.html) which fully 835 | supported only on UNIX. 836 | 837 | Arguments: 838 | func: Function to decorate 839 | time_limit: Max time (in seconds) for function to run, 0 means no time limit 840 | signal_handler: Custom signal handler raising a TimeoutError 841 | signum: Signal number to be used, default=signal.SIGALRM (14) 842 | 843 | Returns: 844 | Decorated function 845 | 846 | Raises: 847 | ValueError: If `time_limit` is not a positive number 848 | TypeError: If `signum` is not an int or an Enum, or if `signal_handler` is not a callable 849 | TimeoutError: If `time_limit` is reached without decorated function finishing 850 | 851 | Usage: 852 | ```python 853 | import time 854 | from deczoo import timeout 855 | 856 | @timeout(time_limit=3) 857 | def add(a, b): 858 | time.sleep(2) 859 | return a+b 860 | 861 | add(1, 2) 862 | 3 863 | 864 | @timeout(time_limit=1) 865 | def add(a, b): 866 | time.sleep(2) 867 | return a+b 868 | 869 | add(1, 2) 870 | > TimeoutError: Reached time limit, terminating add 871 | ``` 872 | """ 873 | 874 | if (not isinstance(time_limit, int)) or time_limit < 0: 875 | raise ValueError("`time_limit` should be a strictly positive number") 876 | 877 | if not isinstance(signum, (int, Enum)): 878 | raise TypeError("`signum` should be an int or an Enum") 879 | 880 | if signal_handler is None: 881 | 882 | def signal_handler(signum, frame): 883 | raise TimeoutError(f"Reached time limit, terminating {func.__name__}") 884 | 885 | signal.signal(signum, signal_handler) # type: ignore 886 | 887 | elif not callable(signal_handler): 888 | raise TypeError("`signal_handler` should be a callable") 889 | 890 | else: 891 | # custom signal handler provided -> bind it to the signal 892 | signal.signal(signum, signal_handler) # type: ignore 893 | 894 | @wraps(func) # type: ignore 895 | def wrapper(*args: PS.args, **kwargs: PS.kwargs) -> R: 896 | signal.alarm(time_limit) 897 | 898 | try: 899 | return func(*args, **kwargs) # type: ignore 900 | except TimeoutError as e: 901 | raise e 902 | except Exception as e: 903 | raise e 904 | finally: 905 | signal.alarm(0) 906 | 907 | return wrapper 908 | 909 | 910 | def raise_if( 911 | condition: Callable[[], bool], 912 | exception: Type[Exception] = Exception, 913 | message: str = "Condition is not satisfied", 914 | ) -> Callable[[Callable[PS, R]], Callable[PS, R]]: 915 | """Raises an exception if `condition` is satisfied. 916 | 917 | Arguments: 918 | condition: Condition to be satisfied 919 | exception: Exception to raise 920 | message: Exception message 921 | 922 | Returns: 923 | Decorated function 924 | 925 | Usage: 926 | ```python 927 | import os 928 | from deczoo import raise_if 929 | 930 | def is_prod(): 931 | '''Returns True if environment is prod''' 932 | return os.environ.get("ENV", "").lower() in {"prod", "production"} 933 | 934 | @raise_if(condition=is_prod, message="Do not run in production!") 935 | def add(a, b): 936 | '''Adds two numbers''' 937 | return a+b 938 | 939 | os.environ["ENV"] = "staging" 940 | add(1, 2) 941 | 3 942 | 943 | os.environ["ENV"] = "prod" 944 | add(1, 2) 945 | > Exception: Do not run in production! 946 | ``` 947 | """ 948 | 949 | def decorator(func: Callable[PS, R]) -> Callable[PS, R]: 950 | @wraps(func) 951 | def wrapper(*args: PS.args, **kwargs: PS.kwargs) -> R: 952 | if condition(): 953 | raise exception(message) 954 | return func(*args, **kwargs) 955 | 956 | return wrapper 957 | 958 | return decorator 959 | --------------------------------------------------------------------------------