inject_spinner decorator

.github/workflows/lint_test_coverage.yml

@@ -117,7 +117,7 @@ jobs:
       - name: install-dependencies
         if: "contains(matrix.python-version, 'pypy-3.10')"
         run: |
-          pip install termcolor==2.3.0 pytest==8.1.2 pytest-xdist==3.6.1
+          pip install termcolor==2.3.0 pytest==8.1.2 pytest-xdist==3.6.1 pytest-mock==3.14.0
 
       - name: run-tests
         if: "contains(matrix.python-version, 'pypy-3.10')"

examples/inject_spinner.py

@@ -0,0 +1,46 @@
+import signal
+import time
+
+from yaspin import inject_spinner
+from yaspin.core import default_handler, Yaspin
+from yaspin.spinners import Spinners
+
+SLEEP = 2
+
+
+# Basic usage
+@inject_spinner()
+def simple_task(spinner: Yaspin, items: list) -> None:
+    for i, _ in enumerate(items, 1):
+        spinner.text = f"Processing item {i}/{len(items)}"
+        time.sleep(SLEEP)
+    spinner.ok("✓")
+
+
+# With spinner customization
+@inject_spinner(Spinners.dots, text="Processing...", color="green")
+def process_items(spinner: Yaspin, items: list) -> None:
+    for i, _ in enumerate(items, 1):
+        spinner.text = f"Processing item {i}/{len(items)}"
+        spinner.color = "yellow" if i % 2 == 0 else "cyan"
+        time.sleep(SLEEP)
+    spinner.ok("✓")
+
+
+# With signal handling
+@inject_spinner(sigmap={signal.SIGINT: default_handler})
+def safe_process(spinner: Yaspin, items: list) -> None:
+    spinner.text = "Press Ctrl+C to interrupt..."
+    time.sleep(5)
+    spinner.ok("✓")
+
+
+def main():
+    items = ["item1", "item2", "item3"]
+    simple_task(items)
+    process_items(items)
+    safe_process(items)
+
+
+if __name__ == "__main__":
+    main()

pyproject.toml

@@ -52,18 +52,20 @@ classifiers = [
 python = "^3.9"
 termcolor = ">=2.2.0, <2.4.0"
 
-[tool.poetry.dev-dependencies]
+[tool.poetry.group.dev.dependencies]
 mypy = "^1.11"
 pytest = "^8.1, <8.2"
 pytest-xdist = "^3.5, !=3.6.0"
 pytest-cov = "^5.0"
+pytest-mock = "^3.14.0"
 semgrep = "^1.89"
 ruff = "^0.7.0"
 
 [tool.poetry.urls]
 "Bug Tracker" = "https://github.com/pavdmyt/yaspin/issues"
 "Changelog" = "https://github.com/pavdmyt/yaspin/releases"
 
+
 [build-system]
 requires = ["poetry_core>=1.0.0"]
 build-backend = "poetry.core.masonry.api"

tests/test_inject_spinner.py

@@ -0,0 +1,203 @@
+from collections.abc import Generator
+from contextlib import contextmanager
+from typing import Any
+
+import signal
+
+import pytest
+
+from yaspin import inject_spinner
+from yaspin.core import Spinner, Yaspin
+from yaspin.spinners import Spinners
+
+
+@pytest.fixture
+def simple_spinner() -> Spinner:
+    return Spinner(["A", "B", "C"], 100)
+
+
+@contextmanager
+def assert_no_yaspin_errors() -> Generator[None, None, None]:
+    """Helper context manager to ensure yaspin starts and stops correctly."""
+    try:
+        yield
+    finally:
+        # Ensure cursor is visible after test
+        print("\033[?25h", end="", flush=True)
+
+
+# Basic functionality tests
+#
+def test_basic_injection() -> None:
+    """Test basic spinner injection without arguments."""
+
+    @inject_spinner()
+    def sample_func(spinner: Yaspin, x: int) -> int:
+        assert isinstance(spinner, Yaspin)
+        return x * 2
+
+    with assert_no_yaspin_errors():
+        result = sample_func(5)
+        assert result == 10
+
+
+def test_spinner_customization(simple_spinner: Spinner) -> None:
+    """Test spinner injection with custom spinner and parameters."""
+
+    @inject_spinner(simple_spinner, text="Testing", color="green")
+    def sample_func(spinner: Yaspin) -> None:
+        assert isinstance(spinner, Yaspin)
+        assert spinner.color == "green"
+        assert spinner.text == "Testing"
+        assert spinner.spinner.frames == ["A", "B", "C"]
+        assert spinner.spinner.interval == 100
+
+    with assert_no_yaspin_errors():
+        sample_func()
+
+
+def test_spinner_property_modification() -> None:
+    """Test modifying spinner properties inside the decorated function."""
+
+    @inject_spinner(text="Initial")
+    def sample_func(spinner: Yaspin) -> None:
+        assert spinner.text == "Initial"
+        spinner.text = "Modified"
+        spinner.color = "red"
+        assert spinner.text == "Modified"
+        assert spinner.color == "red"
+
+    with assert_no_yaspin_errors():
+        sample_func()
+
+
+# Error handling tests
+#
+def test_error_propagation() -> None:
+    """Test error propagation through the decorated function."""
+
+    @inject_spinner()
+    def failing_func(spinner: Yaspin) -> None:
+        raise ValueError("Test error")
+
+    with assert_no_yaspin_errors(), pytest.raises(ValueError, match="Test error"):
+        failing_func()
+
+
+def test_invalid_color() -> None:
+    """Test handling of invalid color parameter."""
+
+    @inject_spinner(color="not_a_color")
+    def invalid_color_func(spinner: Yaspin) -> None:
+        pass
+
+    with pytest.raises(ValueError):
+        invalid_color_func()
+
+
+# Parameter handling tests
+#
+def test_args_kwargs_handling() -> None:
+    """Test handling of various argument combinations."""
+
+    @inject_spinner(Spinners.dots, text="Testing")
+    def sample_func(spinner: Yaspin, *args: Any, **kwargs: Any) -> tuple[tuple, dict]:
+        return args, kwargs
+
+    with assert_no_yaspin_errors():
+        args, kwargs = sample_func(1, 2, a=3, b=4)
+        assert args == (1, 2)
+        assert kwargs == {"a": 3, "b": 4}
+
+
+def test_return_values() -> None:
+    """Test that return values are properly propagated."""
+
+    @inject_spinner()
+    def func_with_return(spinner: Yaspin, x: int, y: int) -> int:
+        return x + y
+
+    with assert_no_yaspin_errors():
+        result = func_with_return(3, 4)
+        assert result == 7
+
+
+# Advanced feature tests
+#
+def test_signal_handling() -> None:
+    """Test that signal handling works correctly."""
+
+    def handler(signum: int, frame: Any, spinner: Yaspin) -> None:
+        spinner.fail("Interrupted")
+
+    @inject_spinner(sigmap={signal.SIGUSR1: handler})
+    def sample_func(spinner: Yaspin) -> None:
+        # Signal handling setup should work without raising errors
+        pass
+
+    with assert_no_yaspin_errors():
+        sample_func()
+
+
+def test_nested_spinners() -> None:
+    """Test nested spinner decorators."""
+
+    @inject_spinner(text="Outer")
+    def outer_func(outer_spinner: Yaspin) -> None:
+        assert isinstance(outer_spinner, Yaspin)
+        assert outer_spinner.text == "Outer"
+
+        @inject_spinner(text="Inner")
+        def inner_func(inner_spinner: Yaspin) -> None:
+            assert isinstance(inner_spinner, Yaspin)
+            assert inner_spinner.text == "Inner"
+
+        inner_func()
+
+    with assert_no_yaspin_errors():
+        outer_func()
+
+
+# Edge cases tests
+#
+def test_empty_spinner() -> None:
+    """Test handling of empty spinner frames."""
+    empty_spinner = Spinner([], 100)
+
+    @inject_spinner(empty_spinner)
+    def empty_spinner_func(spinner: Yaspin) -> None:
+        # Should use default spinner when empty frames are provided
+        assert len(spinner.spinner.frames) > 0
+
+    with assert_no_yaspin_errors():
+        empty_spinner_func()
+
+
+def test_long_text() -> None:
+    """Test handling of very long text."""
+
+    @inject_spinner(text="a" * 1000)
+    def long_text_func(spinner: Yaspin) -> None:
+        # Should handle long text without crashing
+        assert len(spinner.text) > 0
+
+    with assert_no_yaspin_errors():
+        long_text_func()
+
+
+def test_context_manager_behavior(mocker) -> None:
+    """Test that the context manager behavior works correctly."""
+    mock_yaspin = mocker.patch("yaspin.api.yaspin")
+    mock_spinner = mocker.MagicMock()
+    mock_yaspin.return_value.__enter__.return_value = mock_spinner
+
+    @inject_spinner()
+    def sample_func(spinner: Yaspin) -> None:
+        spinner.text = "Working"
+
+    sample_func()
+
+    # Verify context manager was used correctly
+    mock_yaspin.assert_called_once()
+    mock_yaspin.return_value.__enter__.assert_called_once()
+    mock_yaspin.return_value.__exit__.assert_called_once()

yaspin/__init__.py

@@ -1,6 +1,6 @@
 # :copyright: (c) 2021 by Pavlo Dmytrenko.
 # :license: MIT, see LICENSE for more details.
-from .api import kbi_safe_yaspin, yaspin
+from .api import inject_spinner, kbi_safe_yaspin, yaspin
 from .core import Spinner
 
-__all__ = ("yaspin", "kbi_safe_yaspin", "Spinner")
+__all__ = ("yaspin", "kbi_safe_yaspin", "Spinner", "inject_spinner")

yaspin/api.py

@@ -8,12 +8,15 @@
 This module implements the Yaspin API.
 """
 
-from typing import Any
+from typing import Any, Callable, TypeVar
 
+import functools
 import signal
 
 from .core import default_handler, Yaspin
 
+T = TypeVar("T")
+
 
 def yaspin(*args: Any, **kwargs: Any) -> Yaspin:
     """Display spinner in stdout.
@@ -83,6 +86,31 @@ def foo():
     return Yaspin(*args, **kwargs)
 
 
+def inject_spinner(*args: Any, **kwargs: Any) -> Callable[[Callable[..., T]], Callable[..., T]]:
+    """
+    Decorator that injects a yaspin spinner into the decorated function.
+    The spinner is passed as the first argument to the decorated function.
+
+    Example:
+        @inject_spinner(Spinners.dots, text="Processing...", color="green")
+        def process_data(spinner: Yaspin, data: list) -> None:
+            for i, item in enumerate(data):
+                spinner.text = f"Processing item {i+1}/{len(data)}"
+                # Process item...
+            spinner.ok("✓")
+    """
+
+    def decorator(func: Callable[..., T]) -> Callable[..., T]:
+        @functools.wraps(func)
+        def wrapper(*func_args: Any, **func_kwargs: Any) -> T:
+            with yaspin(*args, **kwargs) as spinner:
+                return func(spinner, *func_args, **func_kwargs)
+
+        return wrapper
+
+    return decorator
+
+
 def kbi_safe_yaspin(*args: Any, **kwargs: Any) -> Yaspin:
     """
     Create a Yaspin instance with a default signal handler for SIGINT.