IPSExtractor/core/csv_utils.py

"""Utilidades para preparar filas y encabezados CSV a partir de JSON extraído.

Extrae y aplana campos no-lista, detecta primera lista de objetos y genera
filas combinando campos repetidos con cada item. Permite:
- filtrar por `allowed_paths`
- agregar columnas agregadas para listas específicas (path base)
"""
from __future__ import annotations

from typing import Any, Dict, List, Optional, Tuple


def flatten_non_list_fields(obj: Any, prefix: str = "") -> Dict[str, Any]:
  flat: Dict[str, Any] = {}
  if isinstance(obj, dict):
    for k, v in obj.items():
      path = f"{prefix}/{k}" if prefix else k
      if isinstance(v, dict):
        flat.update(flatten_non_list_fields(v, path))
      elif isinstance(v, list):
        # omit lists here; handled separately as items
        continue
      else:
        flat[path] = v
  else:
    flat[prefix or "value"] = obj
  return flat


def _flatten_item_dict(it: Dict[str, Any], prefix: str = "") -> Dict[str, Any]:
  """Aplana recursivamente un dict de ítem, usando '/' como separador.

  Ejemplo:
    {"V": {"Unitario": 10}, "Codigo": "X"} -> {"V/Unitario": 10, "Codigo": "X"}
  """
  flat: Dict[str, Any] = {}
  for k, v in it.items():
    base = f"{prefix}{k}" if not prefix else f"{prefix}{k}"
    if isinstance(v, dict):
      flat.update(_flatten_item_dict(v, base + "/"))
    else:
      flat[base] = v
  return flat


def _find_first_items_list(obj: Any, exclude_paths: Optional[set[str]] = None, prefix: str = "") -> Tuple[str, List[Dict[str, Any]] | None]:
  exclude_paths = exclude_paths or set()
  if isinstance(obj, dict):
    for k, v in obj.items():
      path = f"{prefix}/{k}" if prefix else k
      if isinstance(v, list) and v and isinstance(v[0], dict):
        if path not in exclude_paths:
          return path, v
      if isinstance(v, dict):
        sub_key, sub_list = _find_first_items_list(v, exclude_paths, path)
        if sub_list is not None:
          return sub_key, sub_list
  return "", None


def _get_at_path(obj: Any, path: str) -> Any:
  parts = [p for p in path.split('/') if p]
  cur = obj
  for p in parts:
    if isinstance(cur, dict) and p in cur:
      cur = cur[p]
    else:
      return None
  return cur


def aggregate_list_fields(data: Dict[str, Any], aggregate_paths: List[str]) -> Dict[str, Any]:
  aggregated: Dict[str, Any] = {}
  for base_path in aggregate_paths:
    lst = _get_at_path(data, base_path)
    if not (isinstance(lst, list) and lst and isinstance(lst[0], dict)):
      continue
    subkeys = set()
    for it in lst:
      subkeys.update([k for k in it.keys()])
    for sk in subkeys:
      values = []
      for it in lst:
        val = it.get(sk)
        if val is None:
          continue
        values.append(str(val))
      key = f"{base_path}/*/{sk}"
      aggregated[key] = ", ".join(values)
  return aggregated


def rows_from_data(
  data: Dict[str, Any],
  aggregate_list_paths: Optional[List[str]] = None,
  allowed_paths: Optional[List[str]] = None,
) -> Tuple[List[Dict[str, Any]], List[str]]:
  """Construye filas y encabezado para CSV desde un JSON estructurado.

  - aggregate_list_paths: rutas base de listas a agregar como columnas unidas por coma.
  - allowed_paths: si se provee, limita las columnas a estas rutas (siempre incluye 'radicado').
  """
  non_list_full = flatten_non_list_fields(data)
  agg_paths = aggregate_list_paths or []
  non_list = dict(non_list_full)
  if agg_paths:
    aggregated = aggregate_list_fields(data, agg_paths)
    if allowed_paths is not None:
      aggregated = {k: v for k, v in aggregated.items() if k in set(allowed_paths)}
    non_list.update(aggregated)

  if allowed_paths is not None:
    allowed_set = set(allowed_paths)
    kept = {k: v for k, v in non_list.items() if (k in allowed_set or k == 'radicado')}
    non_list = kept

  exclude = set(agg_paths) if agg_paths else set()
  items_key, items = _find_first_items_list(data, exclude_paths=exclude)
  if not items:
    header_nonlist: List[str] = []
    if 'radicado' in non_list:
      header_nonlist.append('radicado')
    header_nonlist.extend([k for k in non_list.keys() if k != 'radicado'])
    return [non_list], header_nonlist

  rows: List[Dict[str, Any]] = []
  header_nonlist: List[str] = []
  if 'radicado' in non_list:
    header_nonlist.append('radicado')
  header_nonlist.extend([k for k in non_list.keys() if k != 'radicado'])

  item_keys = set()
  # allowed subrutas dentro de items, p. ej. "V/Unitario"
  allowed_item_subpaths: Optional[set[str]] = None
  if allowed_paths is not None:
    allowed_item_subpaths = set()
    prefix = f"{items_key}/*/"
    for p in allowed_paths:
      if p.startswith(prefix):
        allowed_item_subpaths.add(p[len(prefix):])

  for it in items:
    flat_it = _flatten_item_dict(it)
    for subpath in flat_it.keys():
      if allowed_item_subpaths is not None and subpath not in allowed_item_subpaths:
        continue
      if '/' in items_key:
        item_keys.add(f"{items_key}/*/{subpath}")
      else:
        item_keys.add(subpath)

  header = header_nonlist + sorted(item_keys)
  for it in items:
    row = dict(non_list)
    flat_it = _flatten_item_dict(it)
    if '/' in items_key:
      for subpath, val in flat_it.items():
        if allowed_item_subpaths is not None and subpath not in allowed_item_subpaths:
          continue
        row[f"{items_key}/*/{subpath}"] = val
    else:
      if allowed_item_subpaths is None:
        row.update(flat_it)
      else:
        for subpath, val in flat_it.items():
          if subpath in allowed_item_subpaths:
            row[subpath] = val
    rows.append(row)
  return rows, header