Skip to content

Commit 1a2c038

Browse files
edan-bainglassedan-bainglass
committed
Refactor entity models out of entity classes
1 parent 9bb780c commit 1a2c038

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

48 files changed

+1126
-985
lines changed

‎docs/source/nitpick-exceptions‎

Lines changed: 9 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -33,8 +33,9 @@ py:meth tempfile.TemporaryDirectory
3333
### AiiDA
3434
py:class ReturnType
3535
py:class SelfType
36-
py:class CollectionType
36+
py:class EntityCollectionType
3737
py:class EntityType
38+
py:class EntityModelType
3839
py:class EntityClsType
3940
py:class ProjectType
4041
py:class BackendEntityType
@@ -51,8 +52,9 @@ py:class UserDefinedType
5152
py:class LinkAnnotateType
5253
py:obj ReturnType
5354
py:obj SelfType
54-
py:obj CollectionType
55+
py:obj EntityCollectionType
5556
py:obj EntityType
57+
py:obj EntityModelType
5658
py:obj BackendEntityType
5759
py:obj BackendNodeType
5860
py:obj TransactionType
@@ -62,23 +64,19 @@ py:class aiida.common.ReturnType
6264
py:class aiida.common.SelfType
6365
py:class aiida.common.lang.MethodType
6466
py:class aiida.cmdline.params._shims.C
65-
py:class aiida.orm.entitites.CollectionType
66-
py:class aiida.orm.entitites.EntityType
67-
py:class aiida.orm.entitites.BackendEntityType
6867
py:class aiida.orm.groups.SelfType
69-
py:class aiida.orm.implementation.entitites.EntityType
7068
py:class aiida.engine.processes.functions.FunctionType
7169
py:class aiida.engine.processes.workchains.workchain.MethodType
7270
py:class aiida.orm.entities.EntityInputModel
73-
py:class aiida.orm.entities.EntityModelType
7471
py:class aiida.orm.entities.EntityType
72+
py:class aiida.orm.entities.EntityModelType
7573
py:class aiida.orm.entities.BackendEntityType
76-
py:class aiida.orm.entities.CollectionType
74+
py:class aiida.orm.entities.EntityCollectionType
7775
py:class aiida.orm.implementation.nodes.BackendNodeType
7876
py:class aiida.orm.implementation.storage_backend.TransactionType
7977
py:class aiida.orm.implementation.entities.EntityType
8078
py:class aiida.orm.nodes.data.enum.EnumType
81-
py:class aiida.orm.nodes.NodeType
79+
py:class aiida.orm.nodes.node.NodeType
8280
py:class aiida.storage.psql_dos.orm.ModelType
8381
py:class aiida.storage.psql_dos.orm.SelfType
8482
py:class aiida.storage.psql_dos.orm.entities.SelfType
@@ -91,18 +89,15 @@ py:class aiida.common.lang.ReturnType
9189
py:obj aiida.common.ReturnType
9290
py:obj aiida.common.SelfType
9391
py:obj aiida.common.lang.ReturnType
94-
py:obj aiida.orm.entitites.CollectionType
95-
py:obj aiida.orm.entitites.EntityType
96-
py:obj aiida.orm.entitites.BackendEntityType
9792
py:obj aiida.orm.groups.SelfType
98-
py:obj aiida.orm.entities.CollectionType
93+
py:obj aiida.orm.entities.EntityCollectionType
9994
py:obj aiida.orm.entities.BackendEntityType
10095
py:obj aiida.orm.entities.EntityType
96+
py:obj aiida.orm.entities.EntityModelType
10197
py:obj aiida.orm.implementation.entities.EntityType
10298
py:obj aiida.orm.implementation.nodes.BackendNodeType
10399
py:obj aiida.orm.implementation.storage_backend.TransactionType
104100
py:obj aiida.orm.nodes.node.NodeType
105-
py:obj aiida.orm.nodes.NodeType
106101
py:obj aiida.storage.psql_dos.orm.ModelType
107102
py:obj aiida.storage.psql_dos.orm.SelfType
108103
py:obj aiida.storage.psql_dos.orm.entities.ModelType

‎pyproject.toml‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -330,7 +330,7 @@ enable_error_code = [
330330
'truthy-bool'
331331
]
332332
plugins = [
333-
'pydantic.mypy',
333+
# 'pydantic.mypy',
334334
'sqlalchemy.ext.mypy.plugin'
335335
]
336336
scripts_are_modules = true
@@ -479,7 +479,7 @@ select = [
479479
# Mark some classes as generic, per https://docs.astral.sh/ruff/settings/#lint_pyflakes_extend-generics
480480
# Needed due to https://github.com/astral-sh/ruff/issues/9298
481481
[tool.ruff.lint.pyflakes]
482-
extend-generics = ["aiida.orm.entities.Entity", "aiida.orm.entities.Collection"]
482+
extend-generics = ["aiida.orm.entities.Entity", "aiida.orm.entities.EntityCollection", "aiida.orm.entities.EntityModel"]
483483

484484
[tool.tox]
485485
legacy_tox_ini = """

‎src/aiida/common/pydantic.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -33,8 +33,8 @@ def MetadataField( # noqa: N802
3333
priority: int = 0,
3434
short_name: str | None = None,
3535
option_cls: t.Any | None = None,
36-
orm_class: type[Entity[t.Any, t.Any]] | str | None = None,
37-
orm_to_model: t.Callable[[Entity[t.Any, t.Any], dict[str, t.Any]], t.Any] | None = None,
36+
orm_class: type[Entity[t.Any, t.Any, t.Any]] | str | None = None,
37+
orm_to_model: t.Callable[[Entity[t.Any, t.Any, t.Any], dict[str, t.Any]], t.Any] | None = None,
3838
model_to_orm: t.Callable[[BaseModel], t.Any] | None = None,
3939
exclude_to_orm: bool = False,
4040
is_attribute: bool = True,

‎src/aiida/orm/__init__.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -45,14 +45,14 @@
4545
'CifData',
4646
'Code',
4747
'CodeEntityLoader',
48-
'Collection',
4948
'Comment',
5049
'Computer',
5150
'ComputerEntityLoader',
5251
'ContainerizedCode',
5352
'Data',
5453
'Dict',
5554
'Entity',
55+
'EntityCollection',
5656
'EntityExtras',
5757
'EntityTypes',
5858
'EnumData',

‎src/aiida/orm/authinfos.py‎

Lines changed: 35 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@
2929
__all__ = ('AuthInfo',)
3030

3131

32-
class AuthInfoCollection(entities.Collection['AuthInfo']):
32+
class AuthInfoCollection(entities.EntityCollection['AuthInfo']):
3333
"""The collection of `AuthInfo` entries."""
3434

3535
@staticmethod
@@ -44,40 +44,44 @@ def delete(self, pk: int) -> None:
4444
self._backend.authinfos.delete(pk)
4545

4646

47-
class AuthInfo(entities.Entity['BackendAuthInfo', AuthInfoCollection]):
47+
class AuthInfoModel(entities.EntityModel):
48+
computer: int = MetadataField(
49+
description='The PK of the computer',
50+
is_attribute=False,
51+
orm_class=Computer,
52+
orm_to_model=lambda auth_info, _: cast(AuthInfo, auth_info).computer.pk,
53+
)
54+
user: int = MetadataField(
55+
description='The PK of the user',
56+
is_attribute=False,
57+
orm_class=User,
58+
orm_to_model=lambda auth_info, _: cast(AuthInfo, auth_info).user.pk,
59+
)
60+
enabled: bool = MetadataField(
61+
True,
62+
description='Whether the instance is enabled',
63+
is_attribute=False,
64+
)
65+
auth_params: Dict[str, Any] = MetadataField(
66+
default_factory=dict,
67+
description='Dictionary of authentication parameters',
68+
is_attribute=False,
69+
)
70+
metadata: Dict[str, Any] = MetadataField(
71+
default_factory=dict,
72+
description='Dictionary of metadata',
73+
is_attribute=False,
74+
)
75+
76+
77+
class AuthInfo(entities.Entity['BackendAuthInfo', AuthInfoCollection, AuthInfoModel]):
4878
"""ORM class that models the authorization information that allows a `User` to connect to a `Computer`."""
4979

80+
Model = AuthInfoModel
81+
5082
_CLS_COLLECTION = AuthInfoCollection
51-
PROPERTY_WORKDIR = 'workdir'
5283

53-
class Model(entities.Entity.Model):
54-
computer: int = MetadataField(
55-
description='The PK of the computer',
56-
is_attribute=False,
57-
orm_class=Computer,
58-
orm_to_model=lambda auth_info, _: cast(AuthInfo, auth_info).computer.pk,
59-
)
60-
user: int = MetadataField(
61-
description='The PK of the user',
62-
is_attribute=False,
63-
orm_class=User,
64-
orm_to_model=lambda auth_info, _: cast(AuthInfo, auth_info).user.pk,
65-
)
66-
enabled: bool = MetadataField(
67-
True,
68-
description='Whether the instance is enabled',
69-
is_attribute=False,
70-
)
71-
auth_params: Dict[str, Any] = MetadataField(
72-
default_factory=dict,
73-
description='Dictionary of authentication parameters',
74-
is_attribute=False,
75-
)
76-
metadata: Dict[str, Any] = MetadataField(
77-
default_factory=dict,
78-
description='Dictionary of metadata',
79-
is_attribute=False,
80-
)
84+
PROPERTY_WORKDIR = 'workdir'
8185

8286
def __init__(
8387
self,

‎src/aiida/orm/comments.py‎

Lines changed: 38 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@
2828
__all__ = ('Comment',)
2929

3030

31-
class CommentCollection(entities.Collection['Comment']):
31+
class CommentCollection(entities.EntityCollection['Comment']):
3232
"""The collection of Comment entries."""
3333

3434
@staticmethod
@@ -65,43 +65,46 @@ def delete_many(self, filters: dict) -> List[int]:
6565
return self._backend.comments.delete_many(filters)
6666

6767

68-
class Comment(entities.Entity['BackendComment', CommentCollection]):
68+
class CommentModel(entities.EntityModel):
69+
uuid: UUID = MetadataField(
70+
description='The UUID of the comment',
71+
is_attribute=False,
72+
exclude_to_orm=True,
73+
)
74+
ctime: datetime = MetadataField(
75+
description='Creation time of the comment',
76+
is_attribute=False,
77+
exclude_to_orm=True,
78+
)
79+
mtime: datetime = MetadataField(
80+
description='Modified time of the comment',
81+
is_attribute=False,
82+
exclude_to_orm=True,
83+
)
84+
node: int = MetadataField(
85+
description='Node PK that the comment is attached to',
86+
is_attribute=False,
87+
orm_class='core.node',
88+
orm_to_model=lambda comment, _: cast(Comment, comment).node.pk,
89+
)
90+
user: int = MetadataField(
91+
description='User PK that created the comment',
92+
is_attribute=False,
93+
orm_class='core.user',
94+
orm_to_model=lambda comment, _: cast(Comment, comment).user.pk,
95+
)
96+
content: str = MetadataField(
97+
description='Content of the comment',
98+
is_attribute=False,
99+
)
100+
101+
102+
class Comment(entities.Entity['BackendComment', CommentCollection, CommentModel]):
69103
"""Base class to map a DbComment that represents a comment attached to a certain Node."""
70104

71-
_CLS_COLLECTION = CommentCollection
105+
Model = CommentModel
72106

73-
class Model(entities.Entity.Model):
74-
uuid: UUID = MetadataField(
75-
description='The UUID of the comment',
76-
is_attribute=False,
77-
exclude_to_orm=True,
78-
)
79-
ctime: datetime = MetadataField(
80-
description='Creation time of the comment',
81-
is_attribute=False,
82-
exclude_to_orm=True,
83-
)
84-
mtime: datetime = MetadataField(
85-
description='Modified time of the comment',
86-
is_attribute=False,
87-
exclude_to_orm=True,
88-
)
89-
node: int = MetadataField(
90-
description='Node PK that the comment is attached to',
91-
is_attribute=False,
92-
orm_class='core.node',
93-
orm_to_model=lambda comment, _: cast(Comment, comment).node.pk,
94-
)
95-
user: int = MetadataField(
96-
description='User PK that created the comment',
97-
is_attribute=False,
98-
orm_class='core.user',
99-
orm_to_model=lambda comment, _: cast(Comment, comment).user.pk,
100-
)
101-
content: str = MetadataField(
102-
description='Content of the comment',
103-
is_attribute=False,
104-
)
107+
_CLS_COLLECTION = CommentCollection
105108

106109
def __init__(
107110
self, node: 'Node', user: 'User', content: Optional[str] = None, backend: Optional['StorageBackend'] = None

‎src/aiida/orm/computers.py‎

Lines changed: 38 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@
3131
__all__ = ('Computer',)
3232

3333

34-
class ComputerCollection(entities.Collection['Computer']):
34+
class ComputerCollection(entities.EntityCollection['Computer']):
3535
"""The collection of Computer entries."""
3636

3737
@staticmethod
@@ -64,9 +64,45 @@ def delete(self, pk: int) -> None:
6464
return self._backend.computers.delete(pk)
6565

6666

67-
class Computer(entities.Entity['BackendComputer', ComputerCollection]):
67+
class ComputerModel(entities.EntityModel):
68+
uuid: UUID = MetadataField(
69+
description='The UUID of the computer',
70+
is_attribute=False,
71+
exclude_to_orm=True,
72+
)
73+
label: str = MetadataField(
74+
description='Label for the computer',
75+
is_attribute=False,
76+
)
77+
description: str = MetadataField(
78+
'',
79+
description='Description of the computer',
80+
is_attribute=False,
81+
)
82+
hostname: str = MetadataField(
83+
description='Hostname of the computer',
84+
is_attribute=False,
85+
)
86+
transport_type: str = MetadataField(
87+
description='Transport type of the computer',
88+
is_attribute=False,
89+
)
90+
scheduler_type: str = MetadataField(
91+
description='Scheduler type of the computer',
92+
is_attribute=False,
93+
)
94+
metadata: Dict[str, Any] = MetadataField(
95+
default_factory=dict,
96+
description='Metadata of the computer',
97+
is_attribute=False,
98+
)
99+
100+
101+
class Computer(entities.Entity['BackendComputer', ComputerCollection, ComputerModel]):
68102
"""Computer entity."""
69103

104+
Model = ComputerModel
105+
70106
_logger = AIIDA_LOGGER.getChild('orm.computers')
71107

72108
PROPERTY_MINIMUM_SCHEDULER_POLL_INTERVAL = 'minimum_scheduler_poll_interval'
@@ -76,39 +112,6 @@ class Computer(entities.Entity['BackendComputer', ComputerCollection]):
76112

77113
_CLS_COLLECTION = ComputerCollection
78114

79-
class Model(entities.Entity.Model):
80-
uuid: UUID = MetadataField(
81-
description='The UUID of the computer',
82-
is_attribute=False,
83-
exclude_to_orm=True,
84-
)
85-
label: str = MetadataField(
86-
description='Label for the computer',
87-
is_attribute=False,
88-
)
89-
description: str = MetadataField(
90-
'',
91-
description='Description of the computer',
92-
is_attribute=False,
93-
)
94-
hostname: str = MetadataField(
95-
description='Hostname of the computer',
96-
is_attribute=False,
97-
)
98-
transport_type: str = MetadataField(
99-
description='Transport type of the computer',
100-
is_attribute=False,
101-
)
102-
scheduler_type: str = MetadataField(
103-
description='Scheduler type of the computer',
104-
is_attribute=False,
105-
)
106-
metadata: Dict[str, Any] = MetadataField(
107-
default_factory=dict,
108-
description='Metadata of the computer',
109-
is_attribute=False,
110-
)
111-
112115
def __init__(
113116
self,
114117
label: Optional[str] = None,

0 commit comments

Comments
 (0)