Serializar clases anidadas en Python

El otro día estaba revisando código en un proyecto tratando de mejorar los tests cuando me encontré con un problema completamente nuevo para mí. ¡Así que para que esto no me pase más hoy os cuento como lo resolví y así me queda escrito para mi yo del futuro!

Serializar clases anidadas en Python

En Python es muy habitual trabajar con clases anidadas, es decir, clases que tienen otras clases como atributos en lugar de primitivos (str, int, list, etc…). Imaginemos ahora una clase Deployment que tiene entre sus atributos otras dos clases llamadas Container y Metadata:

import json

from dataclasses import dataclass
from src.extra import Container, Metadata


@dataclass
class Deployment:
    container: Container
    metadata: Metadata
    replicas: int
    service_account_name: str

Estas clases Container y Metadata a su vez también son clases con sus propios atributos:

@dataclass
class Container:
    image: str
    cpu: str
    memory: str
    name: str
    env: list[EnvVar | EnvVarFromSecret] = field(default_factory=list)
    ports: list[Port] = field(default_factory=list)


@dataclass
class Metadata:
    name: str
    labels: dict[str, str]
    annotations: dict[str, str]

Con estos modelos en mente estaba tratando de mejorar nuestros tests haciendo approval testing de tal forma que pudiésemos comparar la representación completa de la clase Deployment en futuras ejecuciones de los tests. De esta forma, en lugar de comprobar algunos campos concretos (ejemplo de testing muy mejorable a continuación) podemos verificar que todos los campos de la clase Deployment sean los esperados.

Este es el test que quería mejorar:

class TestDeployment:
    def test_service_account_name(self) -> None:
        service_account_name = "another-service-account"
        deployment = (DeploymentBuilder()
            .with_service_account_name(service_account_name)
        )

        expect(deployment.service_account_name).to(equal(expected_name))

Y este es el test que querría tener:

def test_good_deployment_with_service_account(self) -> None:
    deployment = (
        DeploymentBuilder()
        .with_service_account_name("another-service-account")
        .build_good()
    )

    verify(deployment) # De esta forma verificamos la clase completamente.

Sin entrar muy en detalle sobre como funciona por debajo el método verify de la librería ApprovalTests, podemos pensar que se serializa el objeto Deployment y lo guarda en disco para usarlo en la siguiente ejecución.

De cara a poder detectar cambios y que estos se muestren en el test de una forma muy clara, opte por implementar el método __repr__ en la clase Deployment de tal forma que se pudiese serializar el objeto Deployment a un formato JSON.

def __repr__(self) -> str:
    return json.dumps(self.__dict__, indent=2)

Sin embargo, al intentar serializar el objeto Deployment me encontré con el siguiente error:

Error

¿Que está pasando?

En mi cabeza, no sé muy bien por qué la verdad, pensaba que si una clase tenía como propiedad, otra clase con propiedades primarias (str, int, list, etc.) se podía serializar sin problemas. Pero no es así, como se puede ver en esta tabla.

La librería JSON de Python out of the box no sabe como serializar clases que no sean todos sus atributos primitivos.

Por suerte existen múltiples formas de solucionar este problema y yo quiero mostrar la que me parece más sencilla y menos verbosa de todas. Si miramos con detenimiento la documentación del método dumps podemos ver que existe un parámetro llamado default que nos permite pasar una función que se encargue de serializar aquellos objetos que no tengan todas sus propiedades primitivas. Este parámetro está definido de la siguiente forma:

If specified, default should be a function that gets called for objects that can’t otherwise be serialized.

Teniendo este conocimiento podemos crear una lambda que llama al atributo __dict__ de cada objeto a serializar si este no es posible serializarlo, este es un atributo especial de todas las clases de Python que devuelve un diccionario con todos los atributos de la misma.

json.dumps(self.__dict__, indent=2, default=lambda o: o.__dict__)

Con este cambio la serialización ya se produce sin problema:

{
  "container": {
    "image": "any-image",
    "cpu": "any-cpu",
    "memory": "any-memory",
    "name": "any-name",
    "env": [
      {
        "name": "ENV",
        "value": "any-value"
      },
      {
        "name": "secret",
        "secret": {
          "name": "secret",
          "key": "key"
        }
      }
    ],
    "ports": [
      {
        "name": "http",
        "number": 80
      }
    ]
  },
  "metadata": {
    "name": "any-name",
    "labels": {
      "app": "any-app"
    },
    "annotations": {
      "author": "any-author"
    }
  },
  "replicas": 1,
  "service_account_name": "any-service-account-name"
}

Con esta representación si se producen cambios en la clase Deployment estos se mostrarán de una forma muy clara en el test y, por tanto, tenemos unos tests mucho más robustos. Ejemplo de un cambio en la clase Deployment añadiendo un nuevo atributo llamado name:

...
-  "service_account_name": "another-service-account"
+  "service_account_name": "another-service-account",
+  "name": "deployment"
...

Conclusión

Todos los ejemplos mostrados en este post están disponibles en GitHub, para que podáis probarlo por vosotros mismos.

¡Espero que os haya gustado!