refactor: Use ensure_context decorator in Actor class

src/apify/_charging.py

@@ -8,8 +8,6 @@
 
 from pydantic import TypeAdapter
 
-from crawlee._utils.context import ensure_context
-
 from apify._models import (
     ActorRun,
     FlatPricePerMonthActorPricingInfo,
@@ -18,7 +16,7 @@
     PricePerDatasetItemActorPricingInfo,
     PricingModel,
 )
-from apify._utils import docs_group
+from apify._utils import docs_group, ensure_context
 from apify.log import logger
 from apify.storages import Dataset
 
@@ -31,6 +29,8 @@
 
 run_validator = TypeAdapter[ActorRun | None](ActorRun | None)
 
+_ensure_context = ensure_context('active')
+
 
 @docs_group('Charging')
 class ChargingManager(Protocol):
@@ -201,7 +201,7 @@ async def __aexit__(
 
         self.active = False
 
-    @ensure_context
+    @_ensure_context
     async def charge(self, event_name: str, count: int = 1) -> ChargeResult:
         def calculate_chargeable() -> dict[str, int | None]:
             """Calculate the maximum number of events of each type that can be charged within the current budget."""
@@ -291,14 +291,14 @@ def calculate_chargeable() -> dict[str, int | None]:
             chargeable_within_limit=calculate_chargeable(),
         )
 
-    @ensure_context
+    @_ensure_context
     def calculate_total_charged_amount(self) -> Decimal:
         return sum(
             (item.total_charged_amount for item in self._charging_state.values()),
             start=Decimal(),
         )
 
-    @ensure_context
+    @_ensure_context
     def calculate_max_event_charge_count_within_limit(self, event_name: str) -> int | None:
         pricing_info = self._pricing_info.get(event_name)
 
@@ -315,7 +315,7 @@ def calculate_max_event_charge_count_within_limit(self, event_name: str) -> int
         result = (self._max_total_charge_usd - self.calculate_total_charged_amount()) / price
         return max(0, math.floor(result)) if result.is_finite() else None
 
-    @ensure_context
+    @_ensure_context
     def get_pricing_info(self) -> ActorPricingInfo:
         return ActorPricingInfo(
             pricing_model=self._pricing_model,
@@ -328,12 +328,12 @@ def get_pricing_info(self) -> ActorPricingInfo:
             },
         )
 
-    @ensure_context
+    @_ensure_context
     def get_charged_event_count(self, event_name: str) -> int:
         item = self._charging_state.get(event_name)
         return item.charge_count if item is not None else 0
 
-    @ensure_context
+    @_ensure_context
     def get_max_total_charge_usd(self) -> Decimal:
         return self._max_total_charge_usd
 

src/apify/_utils.py

@@ -1,13 +1,48 @@
 from __future__ import annotations
 
 import builtins
+import inspect
 import sys
+from collections.abc import Callable
 from enum import Enum
+from functools import wraps
 from importlib import metadata
-from typing import TYPE_CHECKING, Any, Literal
+from typing import Any, Literal, TypeVar, cast
 
-if TYPE_CHECKING:
-    from collections.abc import Callable
+T = TypeVar('T', bound=Callable[..., Any])
+
+
+def ensure_context(attribute_name: str) -> Callable[[T], T]:
+    """Create a decorator that ensures the context manager is initialized before executing the method.
+
+    The decorator checks if the calling instance has the specified attribute and verifies that it is set to `True`.
+    If the instance is inactive, it raises a `RuntimeError`. Works for both synchronous and asynchronous methods.
+
+    Args:
+        attribute_name: The name of the boolean attribute to check on the instance.
+
+    Returns:
+        A decorator that wraps methods with context checking.
+    """
+
+    def decorator(method: T) -> T:
+        @wraps(method)
+        def sync_wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
+            if not getattr(self, attribute_name, False):
+                raise RuntimeError(f'The {self.__class__.__name__} is not active. Use it within the context.')
+
+            return method(self, *args, **kwargs)
+
+        @wraps(method)
+        async def async_wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
+            if not getattr(self, attribute_name, False):
+                raise RuntimeError(f'The {self.__class__.__name__} is not active. Use it within the async context.')
+
+            return await method(self, *args, **kwargs)
+
+        return cast('T', async_wrapper if inspect.iscoroutinefunction(method) else sync_wrapper)
+
+    return decorator
 
 
 def get_system_info() -> dict:

tests/e2e/conftest.py

@@ -69,7 +69,7 @@ def _prepare_test_env() -> None:
         if hasattr(apify._actor.Actor, '__wrapped__'):
             delattr(apify._actor.Actor, '__wrapped__')
 
-        apify._actor.Actor._is_initialized = False
+        apify._actor.Actor._active = False
 
         # Set the environment variable for the local storage directory to the temporary path.
         monkeypatch.setenv(ApifyEnvVars.LOCAL_STORAGE_DIR, str(tmp_path))

tests/e2e/test_actor_lifecycle.py

@@ -118,24 +118,24 @@ async def test_actor_sequential_contexts(make_actor: MakeActorFunction, run_acto
     async def main() -> None:
         async with Actor as actor:
             actor._exit_process = False
-            assert actor._is_initialized is True
+            assert actor._active is True
 
         # Actor after Actor.
         async with Actor as actor:
             actor._exit_process = False
-            assert actor._is_initialized is True
+            assert actor._active is True
 
         # Actor() after Actor.
         async with Actor(exit_process=False) as actor:
-            assert actor._is_initialized is True
+            assert actor._active is True
 
         # Actor() after Actor().
         async with Actor(exit_process=False) as actor:
-            assert actor._is_initialized is True
+            assert actor._active is True
 
         # Actor after Actor().
         async with Actor as actor:
-            assert actor._is_initialized is True
+            assert actor._active is True
 
     actor = await make_actor(label='actor-sequential-contexts', main_func=main)
     run_result = await run_actor(actor)

tests/integration/conftest.py

@@ -60,7 +60,7 @@ def _prepare_test_env() -> None:
         if hasattr(apify._actor.Actor, '__wrapped__'):
             delattr(apify._actor.Actor, '__wrapped__')
 
-        apify._actor.Actor._is_initialized = False
+        apify._actor.Actor._active = False
 
         # Set the environment variable for the local storage directory to the temporary path.
         monkeypatch.setenv(ApifyEnvVars.LOCAL_STORAGE_DIR, str(tmp_path))

tests/unit/actor/test_actor_lifecycle.py

@@ -68,41 +68,41 @@ async def test_actor_init_instance_manual() -> None:
     """Test that Actor instance can be properly initialized and cleaned up manually."""
     actor = Actor()
     await actor.init()
-    assert actor._is_initialized is True
+    assert actor._active is True
     await actor.exit()
-    assert actor._is_initialized is False
+    assert actor._active is False
 
 
 async def test_actor_init_instance_async_with() -> None:
     """Test that Actor instance can be properly initialized and cleaned up using async context manager."""
     actor = Actor()
     async with actor:
-        assert actor._is_initialized is True
+        assert actor._active is True
 
-    assert actor._is_initialized is False
+    assert actor._active is False
 
 
 async def test_actor_init_class_manual() -> None:
     """Test that Actor class can be properly initialized and cleaned up manually."""
     await Actor.init()
-    assert Actor._is_initialized is True
+    assert Actor._active is True
     await Actor.exit()
-    assert not Actor._is_initialized
+    assert not Actor._active
 
 
 async def test_actor_init_class_async_with() -> None:
     """Test that Actor class can be properly initialized and cleaned up using async context manager."""
     async with Actor:
-        assert Actor._is_initialized is True
+        assert Actor._active is True
 
-    assert not Actor._is_initialized
+    assert not Actor._active
 
 
 async def test_fail_properly_deinitializes_actor(actor: _ActorType) -> None:
     """Test that fail() method properly deinitializes the Actor."""
-    assert actor._is_initialized
+    assert actor._active
     await actor.fail()
-    assert actor._is_initialized is False
+    assert actor._active is False
 
 
 async def test_actor_handles_exceptions_and_cleans_up_properly() -> None:
@@ -111,16 +111,16 @@ async def test_actor_handles_exceptions_and_cleans_up_properly() -> None:
 
     with contextlib.suppress(Exception):
         async with Actor() as actor:
-            assert actor._is_initialized
+            assert actor._active
             raise Exception('Failed')  # noqa: TRY002
 
     assert actor is not None
-    assert actor._is_initialized is False
+    assert actor._active is False
 
 
 async def test_double_init_raises_runtime_error(actor: _ActorType) -> None:
     """Test that attempting to initialize an already initialized Actor raises RuntimeError."""
-    assert actor._is_initialized
+    assert actor._active
     with pytest.raises(RuntimeError):
         await actor.init()
 
@@ -196,7 +196,7 @@ def on_event(event_type: Event) -> Callable:
 
     actor = Actor()
     async with actor:
-        assert actor._is_initialized
+        assert actor._active
         actor.on(Event.PERSIST_STATE, on_event(Event.PERSIST_STATE))
         actor.on(Event.SYSTEM_INFO, on_event(Event.SYSTEM_INFO))
         await asyncio.sleep(1)
@@ -249,12 +249,12 @@ async def test_actor_sequential_contexts(*, first_with_call: bool, second_with_c
     mock = AsyncMock()
     async with Actor(exit_process=False) if first_with_call else Actor as actor:
         await mock()
-        assert actor._is_initialized is True
+        assert actor._active is True
 
     # After exiting the context, new Actor instance can be created without conflicts.
     async with Actor() if second_with_call else Actor as actor:
         await mock()
-        assert actor._is_initialized is True
+        assert actor._active is True
 
     # The mock should have been called twice, once in each context.
     assert mock.call_count == 2

tests/unit/conftest.py

@@ -62,7 +62,7 @@ def prepare_test_env(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> Callabl
     def _prepare_test_env() -> None:
         if hasattr(apify._actor.Actor, '__wrapped__'):
             delattr(apify._actor.Actor, '__wrapped__')
-        apify._actor.Actor._is_initialized = False
+        apify._actor.Actor._active = False
 
         # Set the environment variable for the local storage directory to the temporary path.
         monkeypatch.setenv(ApifyEnvVars.LOCAL_STORAGE_DIR, str(tmp_path))