pumpfun-bonkfun-bot/src/cleanup/modes.py

from cleanup.manager import AccountCleanupManager
from utils.logger import get_logger

logger = get_logger(__name__)


def should_cleanup_after_failure(cleanup_mode) -> bool:
    return cleanup_mode == "on_fail"


def should_cleanup_after_sell(cleanup_mode) -> bool:
    return cleanup_mode == "after_sell"


def should_cleanup_post_session(cleanup_mode) -> bool:
    return cleanup_mode == "post_session"


async def handle_cleanup_after_failure(
        client, wallet, mint, priority_fee_manager, cleanup_mode, cleanup_with_prior_fee, force_burn
    ):
    if should_cleanup_after_failure(cleanup_mode):
        logger.info("[Cleanup] Triggered by failed buy transaction.")
        manager = AccountCleanupManager(client, wallet, priority_fee_manager, cleanup_with_prior_fee, force_burn)
        await manager.cleanup_ata(mint)

async def handle_cleanup_after_sell(
        client, wallet, mint, priority_fee_manager, cleanup_mode, cleanup_with_prior_fee, force_burn
    ):
    if should_cleanup_after_sell(cleanup_mode):
        logger.info("[Cleanup] Triggered after token sell.")
        manager = AccountCleanupManager(client, wallet, priority_fee_manager, cleanup_with_prior_fee, force_burn)
        await manager.cleanup_ata(mint)

async def handle_cleanup_post_session(
        client, wallet, mints, priority_fee_manager, cleanup_mode, cleanup_with_prior_fee, force_burn
    ):
    if should_cleanup_post_session(cleanup_mode):
        logger.info("[Cleanup] Triggered post trading session.")
        manager = AccountCleanupManager(client, wallet, priority_fee_manager, cleanup_with_prior_fee, force_burn)
        for mint in mints:
            await manager.cleanup_ata(mint)