20/07/2016 - por Gabriel Marcondes

dict, testes e métodos especiais em Python

Vamos fazer um bem-bolado de Python: inspirado pela explicação dada no excelente livro do Luciano Ramalho, Fluent Python, vamos implementar uma classe similar ao dict do Python, com a filosofia de TDD (test-driven development) e aprendendo como o interpretador usa os métodos especiais.

testes

No desenvolvimento orientado a testes, a gente escreve os testes antes de escrever o código (ou deveria, né?). O ciclo é: escreve testes, roda testes; se os testes falharem escreve código, se passarem escreve mais testes. É importante, e óbvio, que os testes cubram todas as funcionalidades esperadas do futuro código, incluindo os casos de erro.

Não vamos usar nenhum framework de testes aqui, vamos apenas usar o assert padrão do Python, assim não teremos dependência nenhuma pra rodar nosso código. Para se aprofundar mais em testes, recomendo que estude algo como o py.test.

Comecemos com alguns testes bem simples.

In [1]:
# testes
d = MeuDicionario()
d["chave"] = "valor"
assert d["chave"] == "valor"
-----------------------------------------------------
NameError           Traceback (most recent call last)
<ipython-input-1-7c08b2bb1319> in <module>()
      1 # testes
----> 2 d = MeuDicionario()
      3 d["chave"] = "valor"
      4 assert d["chave"] == "valor"

NameError: name 'MeuDicionario' is not defined

Claro que o teste iria falhar, pois sequer implementamos a classe. E se o teste falhou, vamos escrever código.

In [2]:
# código
class MeuDicionario:
    pass

# testes
d = MeuDicionario()
d["chave"] = "valor"
assert d["chave"] == "valor"
-----------------------------------------------------
TypeError           Traceback (most recent call last)
<ipython-input-2-a8d4ecf80135> in <module>()
      5 # testes
      6 d = MeuDicionario()
----> 7 d["chave"] = "valor"
      8 assert d["chave"] == "valor"

TypeError: 'MeuDicionario' object does not support item assignment

Agora o erro acontece na segunda linha dos testes, afinal não implementamos "item assignment" na nossa classe. Já deu pra sacar qual é a do TDD, né? Vamos falar um pouco de como funciona um dicionário, agora.

hash tables and hash maps

O tipo de estrutura de dados que precisamos aprender primeiro é o hash table. Basicamente é um grande array vazio, onde vamos inserindo dados em posições calculadas por uma função hash (uma função que "espalha" valores) a partir do dado original. Assim, quando queremos achar se um valor está nessa tabela, não precisamos procurar por todo o array: calculamos a posição esperada e verificamos a presença já sabendo onde ele estará. O ganho é termos um tempo de pesquisa constante, independente do tamanho dos dados, em troca de gastarmos mais espaço na memória para guardá-los.

Baseado no hash table podemos construir o hash map, que é o uso de uma hash table para guardarmos uma estrutura de mapeamento.

Então isso será o MeuDicionario: grandes arrays vazios, onde vamos inserir os pares chave-valor na posição calculada pela função hash do Python.

métodos especiais

Outra coisa que precisamos saber é como o Python procura por código para os seus operadores. Pela mensagem de erro do último teste, podemos ver que o operador [] procura pela funcionalidade de item assignment. Esta e outras funcionalidades relacionadas aos operadores da linguagem são implementadas nos métodos especiais, métodos com nomes específicos, começados e terminados com dois underlines. Por exemplo, o operador de soma + irá procurar pelo método __add__.

Você pode ver aqui com mais detalhes a explicação de cada método especial e quando cada um é chamado. No nosso caso, queremos o operador de item assignment, então implementaremos o __setitem__.

Vamos começar inicializando nosso grande array vazio e implementando o método __setitem__.

In [3]:
class MeuDicionario:
    def __init__(self):
        self._tamanho = 128
        self._chaves = [None for i in range(128)]
        self._valores = [None for i in range(128)]
    
    def __setitem__(self, chave, valor):
        posicao = hash(chave) % self._tamanho
        self._chaves[posicao] = chave
        self._valores[posicao] = valor

# testes
d = MeuDicionario()
d["chave"] = "valor"
assert d["chave"] == "valor"
-----------------------------------------------------
TypeError           Traceback (most recent call last)
<ipython-input-3-af90d4a998a6> in <module>()
     13 d = MeuDicionario()
     14 d["chave"] = "valor"
---> 15 assert d["chave"] == "valor"

TypeError: 'MeuDicionario' object is not subscriptable

Opa! Passamos pelo teste de inserção, mas falhamos na comparação porque não implementamos ainda o método __getitem__.

Agora que já entendemos como testamos e como implementamos, vamos expandir os testes para mais casos de uso do nosso dicionário. Vamos querer os métodos de listar as chaves e os valores, de verificação de conteúdo e de iteração de itens. Claro, eles vão falhar.

In [4]:
class MeuDicionario:
    def __init__(self):
        self._tamanho = 128
        self._chaves = [None for i in range(128)]
        self._valores = [None for i in range(128)]
    
    def __setitem__(self, chave, valor):
        posicao = hash(chave) % self._tamanho
        self._chaves[posicao] = chave
        self._valores[posicao] = valor

# testes
d = MeuDicionario()
d["chave"] = "valor"
assert d["chave"] == "valor"
assert "chave" in d
assert "valor" in d.values()
for c, v in d.items():
    assert c == "chave"
    assert v == "valor"
-----------------------------------------------------
TypeError           Traceback (most recent call last)
<ipython-input-4-2f0b5c5bfbe1> in <module>()
     13 d = MeuDicionario()
     14 d["chave"] = "valor"
---> 15 assert d["chave"] == "valor"
     16 assert "chave" in d
     17 assert "valor" in d.values()

TypeError: 'MeuDicionario' object is not subscriptable

Os métodos que precisamos são __getitem__, __contains__, values, items.

In [5]:
class MeuDicionario:
    def __init__(self):
        self._tamanho = 128
        self._chaves = [None for i in range(128)]
        self._valores = [None for i in range(128)]
    
    def __setitem__(self, chave, valor):
        posicao = hash(chave) % self._tamanho
        self._chaves[posicao] = chave
        self._valores[posicao] = valor
    
    def __getitem__(self, chave):
        posicao = hash(chave) % self._tamanho
        return self._valores[posicao]
    
    def __contains__(self, chave):
        posicao = hash(chave) % self._tamanho
        return self._chaves[posicao] is not None
    
    def values(self):
        return [valor for valor in self._valores if valor is not None]

    def items(self):
        for par in zip(self._chaves, self._valores):
            if par[0] is not None:
                yield par
        
# testes
d = MeuDicionario()
d["chave"] = "valor"
assert d["chave"] == "valor"
assert "chave" in d
assert "valor" in d.values()
for c, v in d.items():
    assert c == "chave"
    assert v == "valor"

print("sucesso!")
sucesso!

Deu tudo certo! Passamos em todos os testes, mas cá entre nós, nossos testes eram muito simples. Existem mais sutilezas nos dicionários, vamos discutir algumas delas.

Um problema que pode acontecer é a colisão de hash. Quanto mais dados, maior a probabilidade que duas chaves caiam no mesmo hash, porque o array vai enchendo. Então vamos implementar duas coisas novas: aumentar o tamanho do array sempre que ele estiver 30% ocupado, redistribuir os pares depois da expansão, e guardarmos listas de chaves em cada posição, para o caso da colisão ainda acontecer.

Para conseguirmos atingir esses cenários, vamos primeiro testar o caso de colisão com um espaço de array reduzido e sem a expansão automática.

In [6]:
class MeuDicionario:
    def __init__(self, tamanho_inicial=128):
        self._tamanho = tamanho_inicial
        self._chaves = [[] for i in range(self._tamanho)]
        self._valores = [[] for i in range(self._tamanho)]
    
    def __setitem__(self, chave, valor):
        posicao = hash(chave) % self._tamanho
        self._chaves[posicao].append(chave)
        self._valores[posicao].append(valor)
    
    def __getitem__(self, chave):
        posicao = hash(chave) % self._tamanho
        sub_posicao = self._chaves[posicao].index(chave)
        return self._valores[posicao][sub_posicao]
    
    def __contains__(self, chave):
        posicao = hash(chave) % self._tamanho
        return chave in self._chaves[posicao]
    
    def values(self):
        for valores in self._valores:
            if valores is not None:
                for valor in valores:
                    yield valor

    def items(self):
        for par in zip(self._chaves, self._valores):
            if par[0]:
                for chave, valor in zip(par[0], par[1]):
                    yield chave, valor
        
# testes
d = MeuDicionario(tamanho_inicial=1)
d["chave"] = "valor"
assert d["chave"] == "valor"
assert "chave" in d
assert "valor" in d.values()
for c, v in d.items():
    assert c == "chave"
    assert v == "valor"

d["outra chave"] = "outro valor"
assert d["chave"] == "valor"
assert d["outra chave"] == "outro valor"
print("chaves", d._chaves)
print("valores", d._valores)

print("sucesso!")
chaves [['chave', 'outra chave']]
valores [['valor', 'outro valor']]
sucesso!

Imprimimos os arrays de chaves e valores para verificarmos que a colisão aconteceu, e mesmo assim os pares foram preservados, já que os testes passaram. Vamos agora implementar a expansão automática. Cada vez que formos inserir um novo item, vamos verificar se o array já está mais de 30% ocupado, e se for o caso, aumentaremos em 50% o tamanho do array, retirando todos os itens e recolocando-os nas suas novas posições.

In [7]:
class MeuDicionario:
    def __init__(self, tamanho_inicial=128):
        self._tamanho = tamanho_inicial
        self._ocupacao = 0
        self._chaves = [[] for i in range(self._tamanho)]
        self._valores = [[] for i in range(self._tamanho)]
        
    def _expandir(self):
        print("expandindo!")
        pares_existentes = list(self.items())
        self._tamanho = int(1.5 * self._tamanho)
        self._chaves = [[] for i in range(self._tamanho)]
        self._valores = [[] for i in range(self._tamanho)]
        for chave, valor in pares_existentes:
            self[chave] = valor
    
    def __setitem__(self, chave, valor):
        if self._ocupacao / self._tamanho > 0.3:
            self._expandir()        
        posicao = hash(chave) % self._tamanho
        self._chaves[posicao].append(chave)
        self._valores[posicao].append(valor)
        self._ocupacao += 1
    
    def __getitem__(self, chave):
        posicao = hash(chave) % self._tamanho
        sub_posicao = self._chaves[posicao].index(chave)
        return self._valores[posicao][sub_posicao]
    
    def __contains__(self, chave):
        posicao = hash(chave) % self._tamanho
        return chave in self._chaves[posicao]
    
    def values(self):
        for valores in self._valores:
            if valores is not None:
                for valor in valores:
                    yield valor

    def items(self):
        for par in zip(self._chaves, self._valores):
            if par[0]:
                for chave, valor in zip(par[0], par[1]):
                    yield chave, valor
        
# testes
d = MeuDicionario(tamanho_inicial=3)
d["chave"] = "valor"
assert d["chave"] == "valor"
assert "chave" in d
assert "valor" in d.values()
for c, v in d.items():
    assert c == "chave"
    assert v == "valor"

d["outra chave"] = "outro valor"
assert d["chave"] == "valor"
assert d["outra chave"] == "outro valor"
print("chaves", d._chaves)
print("valores", d._valores)
assert d._tamanho > 3

print("sucesso!")
expandindo!
chaves [['chave', 'outra chave'], [], [], []]
valores [['valor', 'outro valor'], [], [], []]
sucesso!

Agora que já temos o dicionário guardando e buscando dados, tratando colisões e se expandindo corretamente, vamos apagar os prints e tratar de um caso errado, buscar uma chave não existente.

In [8]:
class MeuDicionario:
    def __init__(self, tamanho_inicial=128):
        self._tamanho = tamanho_inicial
        self._ocupacao = 0
        self._chaves = [[] for i in range(self._tamanho)]
        self._valores = [[] for i in range(self._tamanho)]
        
    def _expandir(self):
        pares_existentes = list(self.items())
        self._tamanho = int(1.5 * self._tamanho)
        self._chaves = [[] for i in range(self._tamanho)]
        self._valores = [[] for i in range(self._tamanho)]
        for chave, valor in pares_existentes:
            self[chave] = valor
    
    def __setitem__(self, chave, valor):
        if self._ocupacao / self._tamanho > 0.3:
            self._expandir()        
        posicao = hash(chave) % self._tamanho
        self._chaves[posicao].append(chave)
        self._valores[posicao].append(valor)
        self._ocupacao += 1
    
    def __getitem__(self, chave):
        posicao = hash(chave) % self._tamanho
        try:
            sub_posicao = self._chaves[posicao].index(chave)
        except ValueError:
            raise KeyError("Chave não Encontrada")
        return self._valores[posicao][sub_posicao]
    
    def __contains__(self, chave):
        posicao = hash(chave) % self._tamanho
        return chave in self._chaves[posicao]
    
    def values(self):
        for valores in self._valores:
            if valores is not None:
                for valor in valores:
                    yield valor

    def items(self):
        for par in zip(self._chaves, self._valores):
            if par[0]:
                for chave, valor in zip(par[0], par[1]):
                    yield chave, valor
        
# testes
d = MeuDicionario(tamanho_inicial=3)
d["chave"] = "valor"
assert d["chave"] == "valor"
assert "chave" in d
assert "valor" in d.values()
for c, v in d.items():
    assert c == "chave"
    assert v == "valor"

d["outra chave"] = "outro valor"
assert d["chave"] == "valor"
assert d["outra chave"] == "outro valor"
assert d._tamanho > 3

try:
    d["chave inexistente"]
except KeyError:
    pass  # exceção esperada!
except Exception:
    raise  # outra exceção, inesperada

print("sucesso!")
sucesso!

finalizando

O dict do Python é muito mais complexo do que vimos aqui, claro. Se quiser exercitar, estude a documentação oficial e implemente os outros métodos que não fizemos. Os exemplos de uso da documentação podem ser novos casos de testes. Você pode começar com os métodos de impressão, por exemplo, para que tenhamos uma saída inteligível para o print(d).

Os testes, aliás, são nossos melhores amigos. Graças a eles pudemos desenvolver as novas funcionalidades, que envolviam reescrever alguns dos métodos, e verificar logo em seguida que os casos anteriores continuavam funcionando.

O código e o jupyter notebook usados para este post estão neste repo do github.