Busca de empresas paginada por cursor

"""Busca empresas por filtros e percorre as páginas pelo cursor.

    CNPJIA_KEY=... python busca_paginada.py
    CNPJIA_KEY=... python busca_paginada.py --filtro uf=SC --filtro situacao=ATIVA --max-paginas 3

Cada página traz até 20 empresas no perfil basic. A próxima página é a mesma
chamada, com os MESMOS filtros, mais cursor=<meta.next_cursor>; mudar um filtro
invalida o cursor. meta.next_cursor null = última página. A busca exige plano
pago (no Free responde insufficient_plan) e cobra por empresa retornada: veja
https://cnpj.ia.br/precos.
"""

import argparse
import os
import sys
import time

import requests

API = "https://api.cnpj.ia.br"
ESPERA_MAXIMA = 60

# Órgãos do Poder Executivo Federal ativos no DF (natureza jurídica 1015).
FILTROS_PADRAO = {"uf": "DF", "natureza": "1015", "situacao": "ATIVA"}


def get(sessao, url, params):
    """GET com timeout; em 429/503 espera o Retry-After e repete uma vez."""
    for tentativa in (1, 2):
        resp = sessao.get(url, params=params, timeout=(5, 30))
        if tentativa == 1 and resp.status_code in (429, 503):
            espera = resp.headers.get("Retry-After", "")
            if espera.isdigit() and int(espera) <= ESPERA_MAXIMA:
                print(f"HTTP {resp.status_code}: aguardando {espera}s (Retry-After)", file=sys.stderr)
                time.sleep(int(espera))
                continue
        return resp


def paginas(sessao, filtros, max_paginas):
    """Gera uma página (o JSON inteiro) por chamada, seguindo meta.next_cursor."""
    cursor = None
    for _ in range(max_paginas):
        params = dict(filtros, cursor=cursor) if cursor else filtros
        resp = get(sessao, f"{API}/v1/cnpjs", params)
        corpo = resp.json()
        if not resp.ok:
            erro = corpo.get("error", {})
            sys.exit(
                f"HTTP {resp.status_code}: {erro.get('code')}: {erro.get('message')} "
                f"(request_id {erro.get('request_id')})"
            )
        yield corpo
        cursor = corpo["meta"].get("next_cursor")
        if not cursor:
            return


def main():
    parser = argparse.ArgumentParser(description="Busca paginada na API do cnpj.ia.br.")
    parser.add_argument("--filtro", action="append", metavar="CAMPO=VALOR", help="repita para vários filtros")
    parser.add_argument("--max-paginas", type=int, default=2)
    args = parser.parse_args()

    chave = os.environ.get("CNPJIA_KEY")
    if not chave:
        print("Defina CNPJIA_KEY.", file=sys.stderr)
        sys.exit(2)

    filtros = dict(f.split("=", 1) for f in args.filtro) if args.filtro else FILTROS_PADRAO

    sessao = requests.Session()
    sessao.headers["Authorization"] = f"Bearer {chave}"

    n_paginas = n_empresas = creditos = 0
    total = "?"
    for pagina in paginas(sessao, filtros, args.max_paginas):
        n_paginas += 1
        for empresa in pagina["data"]:
            n_empresas += 1
            print(f"{empresa['cnpj']};{empresa['razao_social']}")
        creditos += pagina["meta"]["credits_charged"]
        if n_paginas == 1:  # a contagem (com teto de exibição) vem na primeira página
            total = pagina["meta"].get("total_count_capped") or "?"

    print(
        f"# páginas: {n_paginas}, empresas: {n_empresas}, total do filtro: {total}, créditos cobrados: {creditos}",
        file=sys.stderr,
    )


if __name__ == "__main__":
    main()