lead-lag-scripts/src/lighter/filter_accounts.py

from decimal import Decimal, InvalidOperation
import argparse
import json
import os

try:
    import ijson
    HAS_IJSON = True
except ImportError:
    HAS_IJSON = False


def parse_available(value):
    try:
        return Decimal(str(value)) if value is not None else Decimal('0')
    except (InvalidOperation, ValueError, TypeError):
        return Decimal('0')


def strip_positions(accounts):
    return [{k: v for k, v in acc.items() if k != 'positions'} if isinstance(acc, dict) else acc for acc in accounts]


def filter_accounts(input_path, min_threshold, max_threshold=None):
    min_dec = Decimal(str(min_threshold))
    max_dec = Decimal(str(max_threshold)) if max_threshold is not None else None
    filtered_accounts = []

    def should_keep(total: Decimal) -> bool:
        return total > min_dec and (max_dec is None or total <= max_dec)

    if HAS_IJSON:
        with open(input_path, 'r', encoding='utf-8') as f:
            for item in ijson.items(f, 'item'):
                resp = item.get('response')
                if not resp:
                    continue
                accounts = resp.get('accounts') or []
                total = sum(parse_available(acc.get('available_balance')) for acc in accounts)
                if should_keep(total):
                    filtered_accounts.extend(strip_positions(accounts))
    else:
        with open(input_path, 'r', encoding='utf-8') as f:
            data = json.load(f)
        for item in data:
            resp = item.get('response')
            if not resp:
                continue
            accounts = resp.get('accounts') or []
            total = sum(parse_available(acc.get('available_balance')) for acc in accounts)
            if should_keep(total):
                filtered_accounts.extend(strip_positions(accounts))

    return filtered_accounts


def main():
    parser = argparse.ArgumentParser(description='Filter accounts by total available_balance range.')
    parser.add_argument(
        '--input',
        default=os.path.join('src', 'lighter', 'accounts_0_100000.json'),
        help='Input JSON file path (default: src/lighter/accounts_0_100000.json)'
    )
    parser.add_argument(
        '--output',
        default=os.path.join('src', 'lighter', 'filtered_accounts.json'),
        help='Output JSON file path (default: src/lighter/filtered_accounts.json)'
    )
    parser.add_argument(
        '--min-threshold',
        type=float,
        default=50000,
        help='Minimum total available_balance to include (default: 50000)'
    )
    parser.add_argument(
        '--max-threshold',
        type=float,
        default=150000,
        help='Maximum total available_balance to include; exclude above this (default: 150000)'
    )
    args = parser.parse_args()

    filtered = filter_accounts(args.input, args.min_threshold, args.max_threshold)
    with open(args.output, 'w', encoding='utf-8') as out:
        json.dump(filtered, out, ensure_ascii=False, indent=2)

    print(f'Filtered {len(filtered)} accounts written to: {args.output}')


if __name__ == '__main__':
    main()