Extract `Model` classes out of ORM entity classes #7093

edan-bainglass · 2025-11-07T08:22:03Z

This PR extracts the nested ORM entity models out to the module level for consistency. It builds on the work of #6255 and #6990. Similar to what was done in #5183 for collections.

The of this PR is as follows (referencing the example below):

It lifts the nested model to the module level and properly names it with the corresponding entity prefix
It adds a Model class attribute as a type alias for the module-level model

2 ensures that much of the existing model mechanics are retained while avoiding the nested model.

Before

class SomeEntity(SomeBase):
   class Model(SomeBase.Model):
      ...

After

class SomeEntityModel(SomeBaseModel):
   ...

class SomeEntity(SomeBase):
   Model: TypeAlias = SomeEntityModel

Why do this?

I found it difficult type hinting for the nested model system. I had hoped that providing the model at the module level will resolve this matter, for example, by allowing the definition of EntityModelType = TypeVar('EntityModelType', bound=EntityModel), instead of binding it to Entity.Model, which seemed to not work with mypy.

It ended up being a bit of a typing hell. However, I am uncertain if it is due to the design or a misunderstanding of typing on my part. Help welcomed @danielhollas 🙏

codecov · 2025-11-07T08:23:53Z

Codecov Report

❌ Patch coverage is 74.14075% with 158 lines in your changes missing coverage. Please review.
✅ Project coverage is 24.35%. Comparing base (5c1b2f4) to head (1a2c038).

Files with missing lines	Patch %	Lines
src/aiida/orm/entities.py	42.19%	37 Missing ⚠️
src/aiida/orm/nodes/data/array/kpoints.py	32.56%	29 Missing ⚠️
src/aiida/orm/nodes/data/array/array.py	42.86%	16 Missing ⚠️
src/aiida/orm/nodes/node.py	70.00%	15 Missing ⚠️
src/aiida/orm/nodes/repository.py	23.08%	10 Missing ⚠️
src/aiida/cmdline/commands/cmd_code.py	0.00%	8 Missing ⚠️
src/aiida/orm/nodes/data/array/trajectory.py	70.38%	8 Missing ⚠️
src/aiida/orm/nodes/data/code/installed.py	56.25%	7 Missing ⚠️
src/aiida/orm/nodes/data/singlefile.py	69.57%	7 Missing ⚠️
src/aiida/cmdline/groups/dynamic.py	0.00%	5 Missing ⚠️
... and 9 more

❗ There is a different number of reports uploaded between BASE (5c1b2f4) and HEAD (1a2c038). Click for more details.

HEAD has 1 upload less than BASE

Flag BASE (5c1b2f4) HEAD (1a2c038)

2 1

Additional details and impacted files

@@             Coverage Diff             @@
##             main    #7093       +/-   ##
===========================================
- Coverage   77.61%   24.35%   -53.25%     
===========================================
  Files         566      566               
  Lines       43546    43720      +174     
===========================================
- Hits        33794    10645    -23149     
- Misses       9752    33075    +23323

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

edan-bainglass · 2025-11-08T09:03:31Z

@danielhollas I'm trying to make this one work, but it seems to require a lot of typing tweaks. Anyhow, it would be easier to understand for a reviewer if and when #6990 is merged. For now, if possible, can you make sense of why the test-install tests are failing? They work locally for me, so unclear.

Update

I force-pushed the PR onto a single commit. Check out the commit to see changes in this PR.

edan-bainglass · 2025-11-08T09:12:08Z

src/aiida/common/typing.py

+if sys.version_info >= (3, 10):
+    from typing import TypeAlias
+else:
+    from typing_extensions import TypeAlias


@danielhollas you do this in #7036 in age_entities.py. Moving it here for consistency.

I've opened #7096 as a better solution for this. I don't think it makes sense to put these common types in our own module.

Well, it makes sense if the try...except dance (as you call it) is required, to avoid excessive dancing 😅 I think in #7096 your claim is that the dance itself is not required.

edan-bainglass · 2025-11-08T09:30:09Z

src/aiida/orm/authinfos.py

        self._backend.authinfos.delete(pk)


+class AuthInfoModel(entities.Entity.Model):


As mentioned in the PR description, I provided the Model attribute (see below) to allow retaining Entity.Model access. However, this required me to define it as a TypeAlias (which is recommended by PEP). But it is unclear to me if this is correct, or rather if this helps with typing. More on this in entities.py.

Which PEP do you mean?

EDIT: In this specific files, if I remove the TypeAlias annotation it still type-checks. So I don't think it's needed here?
(also it seems very weird for it to be TypeAlias, it's a (generic) class?

Sorry, not PEP, but mypy - see https://mypy.readthedocs.io/en/stable/common_issues.html#variables-vs-type-aliases. However, this is not quite correct. The recommendation is that the class variable Model should be explicitly typed as a TypeAlias, as compared to global variables. In any case, I don't care for this solution in general. I don't think TypeAlias is the correct approach, but I'm uncertain about all of this.

I'm not sure if it is good to inherit from a class property reference. However, I encounter strange behavior from pydantic when I expose the model classes directly. I'll paste the error here shortly.

Right, so after switching from Entity.Model to EntityModel, etc. in the inheritance chain Entity -> Node -> Data -> AbstractCode -> InstalledCode, I get the following mypy error

(aiida-dev) (refactor-models) aiida-core > mypy src/aiida/orm/nodes/data/code/ src/aiida/orm/nodes/data/code/installed.py:38: error: INTERNAL ERROR -- Please try using mypy master on GitHub: https://mypy.readthedocs.io/en/stable/common_issues.html#using-a-development-mypy-build Please report a bug at https://github.com/python/mypy/issues version: 1.18.2 Traceback (most recent call last): File "mypy/semanal.py", line 7350, in accept File "mypy/nodes.py", line 1440, in accept File "mypy/semanal.py", line 1775, in visit_class_def File "mypy/semanal.py", line 1991, in analyze_class File "mypy/semanal.py", line 2038, in analyze_class_body_common File "mypy/semanal.py", line 2123, in apply_class_plugin_hooks File "/home/edanb/miniforge3/envs/aiida-dev/lib/python3.10/site-packages/pydantic/mypy.py", line 184, in _pydantic_model_class_maker_callback return transformer.transform() File "/home/edanb/miniforge3/envs/aiida-dev/lib/python3.10/site-packages/pydantic/mypy.py", line 518, in transform fields, class_vars = self.collect_fields_and_class_vars(config, is_root_model) File "/home/edanb/miniforge3/envs/aiida-dev/lib/python3.10/site-packages/pydantic/mypy.py", line 682, in collect_fields_and_class_vars self._api.fail( File "mypy/semanal.py", line 7325, in fail AttributeError: attribute 'column' of 'Context' undefined src/aiida/orm/nodes/data/code/installed.py:38: : note: use --pdb to drop into pdb

I have not yet attempted further debugging.

Notice that this is an apparent issue of the pydantic mypy plugin, which you recently introduced in #7064. Indeed, disabling the plugin removes the error.

Oof, we're getting into some dark magic territory. Maybe if you managed to get a minimal reproducer you could open a bug report against pydantic? Internal error seems like something that's not your fault.

(I don't really care about the plugin, I hoped it would be helpful, but we can remove it if it would conflict with this PR)

Btw: this discussion seems relevant
python/typing#1424

Edit: also this?
python/mypy#11538 (comment)

edan-bainglass · 2025-11-08T09:32:56Z

src/aiida/orm/entities.py

+class Entity(abc.ABC, Generic[BackendEntityType, EntityCollectionType], metaclass=EntityFieldMeta):
+    """An AiiDA entity"""
+
+    Model: TypeAlias = EntityModel


As promised, more on EntityModelType...

Here I initially tried to add an EntityModelType to the Generic list on Entity and wire it throughout Entity, for example, Model: type[EntityModelType] = EntityModel. But this did not yield the desired outcome. Again, it may just be due to a misunderstanding of typing on my part. @danielhollas great if we can discuss this.

edan-bainglass · 2025-11-08T10:12:42Z

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

            except IndexError:
                raise ValueError(
-                    'To understand if it is a metal or insulator, ' 'need more bands than n_band=number_electrons'
+                    'To understand if it is a metal or insulator, need more bands than n_band=number_electrons'


Unrelated formatting (a few more later)

edan-bainglass · 2025-11-08T10:22:19Z

src/aiida/orm/nodes/node.py


    @classmethod
-    def from_model(cls, model: Model) -> Self:  # type: ignore[override]
+    def from_model(cls, model: NodeModel) -> Self:  # type: ignore[override]


Unclear why mypy complains. NodeModel is a subclass of EntityModel. But I imagine here I'm misunderstanding typing/subtyping, etc. @danielhollas what do you think?

danielhollas · 2025-11-09T05:11:53Z

'm trying to make this one work, but it seems to require a lot of typing tweaks.

I'll have a look once I am back but no promises as this seem to touch on the generics / variance topics which are still very murky to me.

Anyhow, it would be easier to understand for a reviewer if and when #6990 is merged.

You can push the branch from #6990 to the origin repo, and then make it a target for this PR.

For now, if possible, can you make sense of why the test-install tests are failing? They work locally for me, so unclear.

The job fails at the setup step, not during the tests. Have you looked into it? The error seems to suggest something is broken here, but I don't have time to dig deeper (still on holidays 😅)

edan-bainglass · 2025-11-09T08:39:56Z

For now, if possible, can you make sense of why the test-install tests are failing? They work locally for me, so unclear.

The job fails at the setup step, not during the tests. Have you looked into it? The error seems to suggest something is broken here, but I don't have time to dig deeper (still on holidays 😅)

It seems to fail at verdi code create core.code.installed --non-interactive --config "${CONFIG}/doubler.yaml", which runs fine locally. No worries. Enjoy your vacay 😎

…ate_model`

edan-bainglass requested review from agoscinski and unkcpz as code owners November 7, 2025 08:22

edan-bainglass force-pushed the refactor-models branch from ada7f96 to 0cc3e45 Compare November 7, 2025 12:38

edan-bainglass marked this pull request as draft November 8, 2025 07:20

edan-bainglass force-pushed the refactor-models branch 2 times, most recently from d7c200e to 44b70b5 Compare November 8, 2025 09:09

edan-bainglass commented Nov 8, 2025

View reviewed changes

edan-bainglass force-pushed the refactor-models branch 3 times, most recently from 90d9c1b to 4a87230 Compare November 8, 2025 10:09

edan-bainglass commented Nov 8, 2025

View reviewed changes

edan-bainglass force-pushed the refactor-models branch 4 times, most recently from 043db37 to b0e1a22 Compare November 8, 2025 12:13

edan-bainglass force-pushed the refactor-models branch from 8a13ea3 to f1068b5 Compare November 9, 2025 11:18

GeigerJ2 added this to aiida-core v2.7.2 Nov 11, 2025

edan-bainglass force-pushed the refactor-models branch from f1068b5 to 3144d5a Compare November 11, 2025 17:13

edan-bainglass added 5 commits November 12, 2025 09:44

Fix ORM pydantic schemas

1ca7bb7

Add Kpoints constructor

5035601

Fix Node.from_model

f042534

Missed cast

fbeb32f

Discard direct orm import from cmdline package

02fed45

edan-bainglass added 22 commits November 12, 2025 09:44

Update regression tests

3435796

Centralize model field collection

f8e724c

Move unhandled node attributes gate to Node class

67073e4

Add attributes to Node constructor

2ae4717

Fix code fields

36da6f6

Fix centralized model field collection

1b5a634

Fix tests

fb47977

Enable ArrayNode numpy array POST payloads via model validator

cec3766

Allow arguments in orm_to_model

aa2ab44

Fix formatting

899ec5a

Fix repository_path default in portable code

ad356fb

Fix iterable array posting

7c4c7b6

Fix model inheritance

5d1563a

Fix from_serialized

23397ff

Rename InputModel and as_input_model to CreateModel and `as_cre…

88aa0cd

…ate_model`

Discard exclude_from_cli

dfd2203

Exclude attributes from ORM

60d4654

Discard unstored from from_serialized

a82c6b7

Discard unstored serialization parameter

afabd93

Always pass kwargs to orm_to_model

e17a9a0

Discard redundant class method in favor of correct model selection

36b4d59

Use is_stored for model selection

9bb780c

edan-bainglass force-pushed the refactor-models branch 2 times, most recently from cd0f29f to 42230ac Compare November 12, 2025 09:57

Refactor entity models out of entity classes

1a2c038

edan-bainglass force-pushed the refactor-models branch from 42230ac to 1a2c038 Compare November 12, 2025 10:13

GeigerJ2 moved this to In progress in aiida-core v2.7.2 Nov 13, 2025

GeigerJ2 assigned edan-bainglass Nov 13, 2025

GeigerJ2 removed this from aiida-core v2.7.2 Nov 21, 2025

GeigerJ2 added this to aiida-core v2.8.0 Nov 21, 2025

		self._backend.authinfos.delete(pk)


		class AuthInfoModel(entities.Entity.Model):

Extract Model classes out of ORM entity classes #7093

Are you sure you want to change the base?

Extract Model classes out of ORM entity classes #7093

Uh oh!

Conversation

edan-bainglass commented Nov 7, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Before

After

Why do this?

Uh oh!

codecov bot commented Nov 7, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

edan-bainglass commented Nov 8, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Update

Uh oh!

edan-bainglass Nov 8, 2025

Choose a reason for hiding this comment

Uh oh!

danielhollas Nov 9, 2025

Choose a reason for hiding this comment

Uh oh!

edan-bainglass Nov 9, 2025

Choose a reason for hiding this comment

Uh oh!

edan-bainglass Nov 8, 2025

Choose a reason for hiding this comment

Uh oh!

danielhollas Nov 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

edan-bainglass Nov 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

edan-bainglass Nov 9, 2025

Choose a reason for hiding this comment

Uh oh!

edan-bainglass Nov 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

danielhollas Nov 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

edan-bainglass Nov 8, 2025

Choose a reason for hiding this comment

Uh oh!

edan-bainglass Nov 8, 2025

Choose a reason for hiding this comment

Uh oh!

edan-bainglass Nov 8, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

danielhollas commented Nov 9, 2025

Uh oh!

edan-bainglass commented Nov 9, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Extract `Model` classes out of ORM entity classes #7093

Extract `Model` classes out of ORM entity classes #7093

edan-bainglass commented Nov 7, 2025 •

edited

Loading

codecov bot commented Nov 7, 2025 •

edited

Loading

edan-bainglass commented Nov 8, 2025 •

edited

Loading

danielhollas Nov 9, 2025 •

edited

Loading

edan-bainglass Nov 9, 2025 •

edited

Loading

edan-bainglass Nov 9, 2025 •

edited

Loading

danielhollas Nov 9, 2025 •

edited

Loading

edan-bainglass Nov 8, 2025 •

edited

Loading