Pydantic: consider overloading attribute models instead

[eimrek]

Currently, the pydantic models for nodes contain

1. a generic attributes Dict:

aiida-core/src/aiida/orm/nodes/node.py

Line 240 in d123ac2

attributes: Optional[Dict[str, Any]] = MetadataField(

2. and any actual node attributes are validated at the "root level", by overloading the Model, e.g. for BandsData:

aiida-core/src/aiida/orm/nodes/data/array/bands.py

Line 216 in d123ac2

class Model(KpointsData.Model):

Some possible drawback of this are

1. conceptually, the pydantic model doesn't reflect the attribute location of the ORM objects (so, the attributes, are at root level).

2. When one serializes a node, you get duplication of the attributes at root level and also in the attributes dict, e.g.

   > n = orm.Int(1)
   > n.serialize()
   {'pk': None,
    'uuid': '2bec1c81-d452-4421-a7dd-7691e8c5133d',
    'node_type': 'data.core.int.Int.',
    ...
    'attributes': {'value': 1}, 
    'value': 1} 

   For integers this does not look too bad, but for nodes with a lot of attributes (e.g. Calcjobs, BandsData), it gets quite messy.

3. (vaguely related) the new restapi GET and POST models are not consistent: one has to POST with attributes at root level, while when you GET a node, the attributes will be nested inside attributes dict.

An alternative solution would be to reflect the attribute structure also in the Pydantic models. So e.g. one could define something a separate AttributesModel that can be overloaded by children:

class Node(Entity['BackendNode', NodeCollection['Node']], metaclass=AbstractNodeMeta):
    ...
    class AttributesModel(BaseModel):
        """Generic attributes with everything allowed."""
        class Config:
            extra = 'allow'  # potentially 'forbid' might also make sense.

    class Model(Entity.Model):
        ...
        attributes: Optional['AttributesModel'] = MetadataField(...)

class StructureData(Node):
    ...
    class AttributesModel(Node.AttributesModel):
        """Overload structure attributes."""
        kinds: List[Dict[str, Any]]
        sites: List[Dict[str, Any]]
        cell: List[List[float]]
        pbc: List[bool]
        ...

@edan-bainglass

[edan-bainglass]

Thanks @eimrek. This is very nice! As stated in person, I am supportive in general, but of course we need to test it carefully to ensure it does not break things. Ideally, I would test it along with #7093, which aims to extract the model system out to avoid the nested class structure, similar to #5183, which did the same for entity collections. As mentioned, this is second to getting the REST API up and running, which depends on the pydantic system. However, I don't see this refactor work affecting it much.

[edan-bainglass]

@eimrek this is now implemented in #6990

[GeigerJ2]

@eimrek this is now implemented in #6990

@edan-bainglass 👀 Shouldn't you be on holidays dude :D

[edan-bainglass]

My holidays are very productive 💪