|
24 | 24 | EventSystemInfoData,
|
25 | 25 | )
|
26 | 26 |
|
| 27 | +from apify._charging import ChargeResult, ChargingManager, ChargingManagerImplementation |
27 | 28 | from apify._configuration import Configuration
|
28 | 29 | from apify._consts import EVENT_LISTENERS_TIMEOUT
|
29 | 30 | from apify._crypto import decrypt_input_secrets, load_private_key
|
@@ -97,6 +98,8 @@ def __init__(
|
97 | 98 | )
|
98 | 99 | )
|
99 | 100 |
|
| 101 | + self._charging_manager = ChargingManagerImplementation(self._configuration, self._apify_client) |
| 102 | + |
100 | 103 | self._is_initialized = False
|
101 | 104 |
|
102 | 105 | @ignore_docs
|
@@ -227,6 +230,10 @@ async def init(self) -> None:
|
227 | 230 | # https://github.com/apify/apify-sdk-python/issues/146
|
228 | 231 |
|
229 | 232 | await self._event_manager.__aenter__()
|
| 233 | + self.log.debug('Event manager initialized') |
| 234 | + |
| 235 | + await self._charging_manager.__aenter__() |
| 236 | + self.log.debug('Charging manager initialized') |
230 | 237 |
|
231 | 238 | self._is_initialized = True
|
232 | 239 | _ActorType._is_any_instance_initialized = True
|
@@ -269,6 +276,7 @@ async def finalize() -> None:
|
269 | 276 | await self._event_manager.wait_for_all_listeners_to_complete(timeout=event_listeners_timeout)
|
270 | 277 |
|
271 | 278 | await self._event_manager.__aexit__(None, None, None)
|
| 279 | + await self._charging_manager.__aexit__(None, None, None) |
272 | 280 |
|
273 | 281 | await asyncio.wait_for(finalize(), cleanup_timeout.total_seconds())
|
274 | 282 | self._is_initialized = False
|
@@ -452,19 +460,46 @@ async def open_request_queue(
|
452 | 460 | storage_client=storage_client,
|
453 | 461 | )
|
454 | 462 |
|
455 |
| - async def push_data(self, data: dict | list[dict]) -> None: |
| 463 | + @overload |
| 464 | + async def push_data(self, data: dict | list[dict]) -> None: ... |
| 465 | + @overload |
| 466 | + async def push_data(self, data: dict | list[dict], charged_event_name: str) -> ChargeResult: ... |
| 467 | + async def push_data(self, data: dict | list[dict], charged_event_name: str | None = None) -> ChargeResult | None: |
456 | 468 | """Store an object or a list of objects to the default dataset of the current Actor run.
|
457 | 469 |
|
458 | 470 | Args:
|
459 | 471 | data: The data to push to the default dataset.
|
| 472 | + charged_event_name: If provided and if the Actor uses the pay-per-event pricing model, |
| 473 | + the method will attempt to charge for the event for each pushed item. |
460 | 474 | """
|
461 | 475 | self._raise_if_not_initialized()
|
462 | 476 |
|
463 | 477 | if not data:
|
464 |
| - return |
| 478 | + return None |
| 479 | + |
| 480 | + data = data if isinstance(data, list) else [data] |
| 481 | + |
| 482 | + max_charged_count = ( |
| 483 | + self._charging_manager.calculate_max_event_charge_count_within_limit(charged_event_name) |
| 484 | + if charged_event_name is not None |
| 485 | + else None |
| 486 | + ) |
465 | 487 |
|
466 | 488 | dataset = await self.open_dataset()
|
467 |
| - await dataset.push_data(data) |
| 489 | + |
| 490 | + if max_charged_count is not None and len(data) > max_charged_count: |
| 491 | + # Push as many items as we can charge for |
| 492 | + await dataset.push_data(data[:max_charged_count]) |
| 493 | + else: |
| 494 | + await dataset.push_data(data) |
| 495 | + |
| 496 | + if charged_event_name: |
| 497 | + return await self._charging_manager.charge( |
| 498 | + event_name=charged_event_name, |
| 499 | + count=min(max_charged_count, len(data)) if max_charged_count is not None else len(data), |
| 500 | + ) |
| 501 | + |
| 502 | + return None |
468 | 503 |
|
469 | 504 | async def get_input(self) -> Any:
|
470 | 505 | """Get the Actor input value from the default key-value store associated with the current Actor run."""
|
@@ -513,6 +548,23 @@ async def set_value(
|
513 | 548 | key_value_store = await self.open_key_value_store()
|
514 | 549 | return await key_value_store.set_value(key, value, content_type=content_type)
|
515 | 550 |
|
| 551 | + def get_charging_manager(self) -> ChargingManager: |
| 552 | + """Retrieve the charging manager to access granular pricing information.""" |
| 553 | + self._raise_if_not_initialized() |
| 554 | + return self._charging_manager |
| 555 | + |
| 556 | + async def charge(self, event_name: str, count: int = 1) -> ChargeResult: |
| 557 | + """Charge for a specified number of events - sub-operations of the Actor. |
| 558 | +
|
| 559 | + This is relevant only for the pay-per-event pricing model. |
| 560 | +
|
| 561 | + Args: |
| 562 | + event_name: Name of the event to be charged for. |
| 563 | + count: Number of events to charge for. |
| 564 | + """ |
| 565 | + self._raise_if_not_initialized() |
| 566 | + return await self._charging_manager.charge(event_name, count) |
| 567 | + |
516 | 568 | @overload
|
517 | 569 | def on(
|
518 | 570 | self, event_name: Literal[Event.PERSIST_STATE], listener: EventListener[EventPersistStateData]
|
|
0 commit comments