eleicoes

eleicoes.py

import os
import csv
import unicodedata
import sqlite3
from os.path import join, isfile
from collections import namedtuple

Linha = namedtuple('Linha', 'cidade,cargo,partido,candidato,zona')

def ler_dataset(caminho):
    colunas = {}

    def c(nome_coluna):
        texto = linha[colunas[nome_coluna]]
        nfkd_form = unicodedata.normalize('NFKD', texto.strip())
        only_ascii = nfkd_form.encode('ASCII', 'ignore')
        return only_ascii.upper().decode()

    with open(caminho, newline='') as csvfile:
        reader = csv.reader(csvfile, delimiter=';', quotechar='"')

        for i, coluna in enumerate(next(reader)):
            colunas[coluna.lower()] = i

        for linha in reader:
            if c('nr_turno') != '1':
                continue

            yield Linha(
                cidade=(c('nm_municipio'), c('sg_uf')),
                cargo=c('ds_cargo'),
                partido=(c('nm_partido'), c('sg_partido')),
                candidato=(c('nm_candidato'), c('nr_candidato'), c('sq_candidato')),
                zona=(c('nr_zona'), c('qt_votos_nominais'))
            )

def criar_tabelas(conn):
    cursor = conn.cursor()

    cursor.execute('''
        create table cidades (
            id int primary key,
            nome varchar(100) not null,
            uf varchar(2) not null
        );''')

    cursor.execute('''
        create table cargos (
            id int primary key,
            nome varchar(100) not null
        );
    ''')

    cursor.execute('''
        create table partidos (
            id int primary key,
            nome varchar(100) not null,
            sigla varchar(10) not null
        );
    ''')

    cursor.execute('''
        create table candidatos (
            id int primary key,
            nome varchar(100) not null,
            numero int  not null,
            partido_id int not null references partidos (id),
            cargo_id int not null references cargos (id)
        );
    ''')

    cursor.execute('''
        create table zonas (
            id int primary key,
            numero int not null,
            cidade_id int not null references cidades (id)
        );
    ''')

    cursor.execute('''
        create table votos (
            zona_id int not null references zonas (id),
            candidato_id int not null references candidatos (id),
            quantidade int not null
        );
    ''')

    conn.commit()

def popular_tabelas(conn):
    cursor = conn.cursor()

    PASTA_BASE = 'datasets'

    id_serial = 1

    cidades = {}
    cargos = {}
    partidos = {}
    candidatos = {}
    zonas = {}

    for nome_arq in os.listdir(PASTA_BASE):
        print('Lendo ' + nome_arq)

        for item in ler_dataset(join(PASTA_BASE, nome_arq)):
            if item.cidade not in cidades:
                nome, uf = item.cidade
                cursor.execute(f"insert into cidades (id, nome, uf) values ({id_serial}, '{nome}', '{uf}')")
                cidades[item.cidade] = id_serial
                id_serial += 1

            if item.cargo not in cargos:
                cursor.execute(f"insert into cargos (id, nome) values ({id_serial}, '{item.cargo}')")
                cargos[item.cargo] = id_serial
                id_serial += 1

            if item.partido[1] not in partidos:
                nome, sigla = item.partido
                cursor.execute(f"insert into partidos (id, nome, sigla) values ({id_serial}, '{nome}', '{sigla}')")
                partidos[sigla] = id_serial
                id_serial += 1

            if item.candidato[2] not in candidatos:
                nome, numero, sequencial = item.candidato
                partido_id = partidos[item.partido[1]]
                cargo_id = cargos[item.cargo]
                cursor.execute(f"insert into candidatos (id, nome, numero, partido_id, cargo_id) values ({id_serial}, '{nome}', {numero}, {partido_id}, {cargo_id})")
                candidatos[sequencial] = id_serial
                id_serial += 1

            if item.zona[0] not in zonas:
                numero = item.zona[0]
                cidade_id = cidades[item.cidade]
                cursor.execute(f"insert into zonas (id, numero, cidade_id) values ({id_serial}, {numero}, {cidade_id})")
                zonas[numero] = id_serial
                id_serial += 1

            quantidade = item.zona[1]
            zona_id = zonas[item.zona[0]]
            candidato_id = candidatos[item.candidato[2]]
            cursor.execute(f"insert into votos (zona_id, candidato_id, quantidade) values ({zona_id}, {candidato_id}, {quantidade})")

    conn.commit()

    print(id_serial)

def main():
    NOME_DB = 'eleicoes-2018.db'
    banco_ja_existe = isfile(NOME_DB)
    conn = sqlite3.connect(NOME_DB)

    try:
        if not banco_ja_existe:
            criar_tabelas(conn)
            popular_tabelas(conn)
    finally:
        conn.close()
        print('fechou')

if __name__ == '__main__':
    main()