Skip to content

Commit

Permalink
Typing: Correct type annotation of WorkChain.on_wait (#5836)
Browse files Browse the repository at this point in the history
It was incorrectly using the `Awaitable` defined by `aiida-core` instead
of the `typing.Awaitable` class, which is what the actual notation is
of the method defined by the `plumpy.processes.Process` base class.

The annotation of `Protect.final` is made more specific in passing.
  • Loading branch information
chrisjsewell committed Jul 6, 2023
1 parent 4db54b7 commit 923cc31
Show file tree
Hide file tree
Showing 2 changed files with 7 additions and 3 deletions.
9 changes: 6 additions & 3 deletions aiida/engine/processes/workchains/workchain.py
Expand Up @@ -43,6 +43,9 @@ class WorkChainSpec(ProcessSpec, PlumpyWorkChainSpec):
pass


MethodType = t.TypeVar('MethodType')


class Protect(ProcessStateMachineMeta):
"""Metaclass that allows protecting class methods from being overridden by subclasses.
Expand Down Expand Up @@ -84,14 +87,14 @@ def __is_final(mcs, method) -> bool: # pylint: disable=unused-private-member
return False

@classmethod
def final(mcs, method: t.Any):
def final(mcs, method: MethodType) -> MethodType:
"""Decorate a method with this method to protect it from being overridden.
Adds the ``__SENTINEL`` object as the ``__final`` private attribute to the given ``method`` and wraps it in
the ``typing.final`` decorator. The latter indicates to typing systems that it cannot be overridden in
subclasses.
"""
method.__final = mcs.__SENTINEL # pylint: disable=protected-access,unused-private-member
method.__final = mcs.__SENTINEL # type: ignore[attr-defined] # pylint: disable=protected-access,unused-private-member
return t.final(method)


Expand Down Expand Up @@ -365,7 +368,7 @@ def on_exiting(self) -> None:
self.logger.exception('exception in _store_nodes called in on_exiting')

@Protect.final
def on_wait(self, awaitables: t.Sequence[Awaitable]):
def on_wait(self, awaitables: t.Sequence[t.Awaitable]):
"""Entering the WAITING state."""
super().on_wait(awaitables)
if self._awaitables:
Expand Down
1 change: 1 addition & 0 deletions docs/source/nitpick-exceptions
Expand Up @@ -58,6 +58,7 @@ py:class aiida.orm.entitites.BackendEntityType
py:class aiida.orm.groups.SelfType
py:class aiida.orm.implementation.entitites.EntityType
py:class aiida.engine.processes.functions.FunctionType
py:class aiida.engine.processes.workchains.workchain.MethodType
py:class aiida.orm.entities.EntityType
py:class aiida.orm.entities.BackendEntityType
py:class aiida.orm.entities.CollectionType
Expand Down

0 comments on commit 923cc31

Please sign in to comment.