sidra-sql/src/sidra_sql/utils.py at main · Quantilica/sidra-sql · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
import itertools
from collections.abc import Iterable

from sidra_fetcher.agregados import (
    Categoria,
    Classificacao,
    Variavel,
)


def unnest_dimensoes(
    variaveis: list[Variavel],
    classificacoes: list[Classificacao],
) -> Iterable[dict]:
    """Expand variables × classification categories into flat Dimensao rows.

    For each variable, computes the cartesian product of all categories
    across every classification (up to 6, mapped to d4–d9).  The unit of
    measure (``mc``/``mn``) is resolved with the following precedence:

    1. The category's own ``unidade`` field (when not ``None``).
    2. The variable's ``unidade`` field as a fallback.

    Args:
        variaveis: Iterable of :class:`~sidra_fetcher.agregados.Variavel`.
        classificacoes: Iterable of
            :class:`~sidra_fetcher.agregados.Classificacao`.

    Returns:
        A list of :class:`~sidra_sql.models.Dimensao` instances,
        one per (variavel, combination-of-categories) tuple.
    """

    # Pre-build a list of (categoria_list,) per classificacao so that
    # itertools.product can expand them correctly.
    cats_per_classificacao: list[list[Categoria]] = [
        classificacao.categorias for classificacao in classificacoes
    ]

    # Pad slots d4–d9: the model supports up to 6 classifications.
    MAX_CLASSIFICACOES = 6

    for variavel in variaveis:
        variavel_id = str(variavel.id)
        variavel_nome = variavel.nome
        unidade_id = None
        unidade_nome = variavel.unidade

        if not cats_per_classificacao:
            # No classifications: yield one row per variable with null d4–d9.
            yield dict(
                mc=unidade_id,
                mn=unidade_nome,
                d2c=variavel_id,
                d2n=variavel_nome,
                d4c=None,
                d4n=None,
                d5c=None,
                d5n=None,
                d6c=None,
                d6n=None,
                d7c=None,
                d7n=None,
                d8c=None,
                d8n=None,
                d9c=None,
                d9n=None,
            )
            continue

        # Cartesian product across all classifications.
        for combo in itertools.product(*cats_per_classificacao):
            # Resolve unit: first category that provides one wins;
            # fall back to the variable's own unit.
            for cat in combo:
                if cat.unidade is not None:
                    unidade_nome = cat.unidade
                    break

            # Map combo slots → d4…d9 (pad with None when fewer than 6).
            padded = list(combo) + [None] * (MAX_CLASSIFICACOES - len(combo))

            def _id(cat):
                return str(cat.id) if cat is not None else None

            def _nome(cat):
                return cat.nome if cat is not None else None

            yield dict(
                mc=unidade_id,
                mn=unidade_nome,
                d2c=variavel_id,
                d2n=variavel_nome,
                d4c=_id(padded[0]),
                d4n=_nome(padded[0]),
                d5c=_id(padded[1]),
                d5n=_nome(padded[1]),
                d6c=_id(padded[2]),
                d6n=_nome(padded[2]),
                d7c=_id(padded[3]),
                d7n=_nome(padded[3]),
                d8c=_id(padded[4]),
                d8n=_nome(padded[4]),
                d9c=_id(padded[5]),
                d9n=_nome(padded[5]),
            )