Maybe AiiDA data type nodes should _always_ be immutable?

[mbercx]

The Problem ❗

The fact that the behaviour for data type nodes is dependent on whether the node is stored or not can lead to confusing and downright erroneous behaviour. An example I encountered recently can be found in #5970. Another discussion that illustrates the problems with mutability when stored or not is #5297.

This concept is also often difficult to explain to new users: "Why is the behaviour of this node suddenly different"? "Why can't I just change an input in this dictionary and run again?". We often have to explain the concept of stored nodes several times before a user is on board, and even then it often leads to errors or frustration, especially when the expected behaviour fails to hold up silently as above.

So I have the following radical suggestion: what if we simply posit that AiiDA nodes are all immutable objects, once and for all?

First thoughts 💭

The idea would be that every data node instance can be fully defined using its constructor (I'm looking at you, KpointsData). Once defined, there are no methods on the node that can change its contents. We basically can follow a str-like approach: it can have methods that sort of sound like they mutate the contents, but these always return a new instance, e.g.:

In [1]: s = 'test'

In [2]: s.capitalize()
Out[2]: 'Test'

In [3]: s.strip('t')
Out[3]: 'es'

We should then consider which methods we want to keep. For StructureData, I would remove almost everything. We don't need to implement methods that are already implemented elsewhere for very similar classes. Again: an AiiDA data node is simply a data container: you put something in it, and need useful methods to obtain the data it stores. Other behaviour might be trickier, the primordial examples to consider are List and Dict, e.g.:

In [1]: d = Dict({'a': {'b': 1, 'c': 2}})

In [2]: d['a'] = 1

We could make such operation raise, or possibly return a new Dict node. I'm currently leaning towards simply not implementing any method that mutates the object, but I'm not sure if I can defend that approach if we start thinking about the consequences in more detail. ^^ Some might argue that we instead were planning to move close to the base Python type (e.g. #2676), but I think we all agree that the AiiDA base types will never exactly replicate the behaviour of the base Python ones. The fact that we semi-advertise the nodes as such again leads to confusion and frustration, so we should move away from that and clearly state that they are not the same. Perhaps even consider renaming them (e.g. DictData, ListData). That may be going too far, but I'm just spit-balling at this point (EDIT: On the other hand, it would make backwards-compatibility easier, see below).

Other considerations 🤔

If we are going to look at redesigning the data nodes, which has already long been on the agenda for the materials-science types, we should consider a few other principles in the design:

1. We should consider making our nodes pydantic BaseModels, and do away with all these *args and **kwargs in the constructors and methods. Although documentation is great, users should really be able to see clearly how to use a certain node type after loading it by looking at the signatures of the methods with some good type argument names, type hinting and docstrings. There may be other reasons not to use pydantic, I'm not set on that, just have been using it a lot lately and now am seeing pydantic-opportunities everywhere. 😅

2. The constructor should have no more arguments than strictly necessary to define the node instance. Notorious examples I'm not a fan of are the ase and pymatgen (and even more egregious, pymatgen_structure and pymatgen_molecule) input arguments of the StructureData constructor. Although being able to convert instances of these classes is essential, it should in my view be done through from_pymatgen()-like classmethods and something like as_pymatgen() to obtain the corresponding pymatgen Structure instance. One aspect of pydantic is that it sort of enforces this. It's also much easier to get good constructor type hinting, especially for classes with a convoluted hierarchy, but I've already raved about this here so I won't repeat myself on that front.

But what about backwards-compatibility? 😱

Of course, changing the way these fundamental data types work is about as backwards-incompatible as you can get, and might require a lot of code changes in plugins. So we'll have to take care in how to approach this. There is some good news though:

1. As I've already mentioned, we were planning to move the "atomistic" data types (a.k.a. materials science data types) out of AiiDA-core into aiida-atomistic (mostly, see Prepare the move of material science specific data classes to separate plugin #2686). That would be an ideal opportunity to make sure all the new versions are immutable at all times. Plugin developers can then move to the new data types at their leisure (but there are still many technical considerations for sure.
2. As for the "core" data types, there are plenty of them that already are immutable. I'd have to have a close look at all them, but Int, Float, Str, Bool obviously come to mind. The codes as well, but I'm actually a bit of a hypocrite on these (see below). Most of the others I'd frankly have to try to be sure, but this post is getting long enough as is.
3. I still have high hopes for aiida-upgrade. You can't fix everything with static analysis, obviously, but if you are a bit rash in making assumptions and let the code developers decide on which automated changes make sense would go a long way towards getting all plugins up to date (even still for v2.X, frankly).

You're a total hypocrite! 🤬

It is true. I have been known to argue that codes and computers should be editable since so many starting users have trouble with this (see aiidateam/team-compass#12). But I just noticed that @ltalirz actually has suggested a very nice comprimise there, so maybe I can have my 🍰 and eat it too.

If you've gotten this far: kudos! For those just skimming: let me sum it up.

TL;DR

1. AiiDA data nodes should always be immutable, eliminating the difference in behaviour when stored or not.
2. We should have as few methods as possible that imitate mutating behaviour, but these should return a new instance with the requested change.
3. Their constructors should be complete but minimal, possibly relying on pydantic BaseModels.
4. Converting from/to AiiDA data types to/from other similar data types in other packages (e.g. StructureData/ase.Atoms/pymatgen.Structure) should be done via from_Xandas_X` methods.

[sphuber]

Nice write-up @mbercx . As you say, this is not an easy topic especially when you have to ensure you cover all bases. I have just given it a read and wanted to quickly write down my first thought before I forget. Likewise these are not complete and are subject to change 😅

I agree that the current design leads to surpising behavior for users, although I really think this mostly applies to Dict and List. Note that similar caveats apply to the Python built-in equivalents. For example, most new Python users will eventually get stumped by weird behavior when they decided to use a dict or list as a function argument default. So I think the quirks surrounding mutability are not unique to AiiDA and the concept of storing alone.

I think, then, that the proposed solution will most likely not be immune to these problems either, and will have quirks and pitfalls of its own that will stump users.
If we were to imagine that your proposal is implemented, what would a user expect in the following example:

d = Dict({'a': 1, 'b': 2})
d['a'] = 5
print(d['a'])

What would this print? A user would probably expect 5 but it would actually print 1. You now have to start explaining to users that if they mutate a node, they have to assign the result to a new variable as a clone is returned. I think this may therefore result in the same level of confusion as before. Note also that there are no warnings here. So a user will have no idea that things weren't working as intended until much later, at which the bughunt will start, similar to yours when you attempted to modify a nested key in a stored Dict node and it being essentially a no-op.

One thing I definitely agree with is that Data classes should provide constructors that allow it to be initialized in a complete state. I would also be tempted to enforce this, except for the fact that I am not sure how and also that there may actually be valid use-cases where this is not impossible and multi-step initialization would be required. It would be frustrating to users that you provide them flexibility for custom use-cases through plugins, but not really.

As for using pydantic: I think it can be helpful, but wouldn't require or enforce this. For example, I am not sure if this is supported, but can you specify attributes to consume files? Or at the very least file-like objects? This is a very common use-case for Data nodes and not sure if pydantic supports it.