Improve typing of QueryBuilder.all() (#6966)

* Provide an overloaded method definitions
for `flat=True` and `flat=False` arguments.
* Bump mypy to 1.17
* Cast away the casts

Author: danielhollas

pyproject.toml

@@ -231,7 +231,7 @@ notebook = [
 ]
 pre-commit = [
   'aiida-core[atomic_tools,rest,tests,tui]',
-  'mypy~=1.16.0',
+  'mypy~=1.17.0',
   'packaging~=23.0',
   'pre-commit~=3.5',
   'sqlalchemy[mypy]~=2.0',

src/aiida/orm/entities.py

@@ -14,7 +14,7 @@
 import pathlib
 from enum import Enum
 from functools import lru_cache
-from typing import TYPE_CHECKING, Any, Generic, List, Optional, Type, TypeVar, Union, cast
+from typing import TYPE_CHECKING, Any, Generic, List, Optional, Type, TypeVar, Union
 
 from plumpy.base.utils import call_with_super_check, super_check
 from pydantic import BaseModel
@@ -156,14 +156,14 @@ def find(
         :return: a list of resulting matches
         """
         query = self.query(filters=filters, order_by=order_by, limit=limit)
-        return cast(List[EntityType], query.all(flat=True))
+        return query.all(flat=True)
 
     def all(self) -> List[EntityType]:
         """Get all entities in this collection.
 
         :return: A list of all entities
         """
-        return cast(List[EntityType], self.query().all(flat=True))
+        return self.query().all(flat=True)
 
     def count(self, filters: Optional['FilterType'] = None) -> int:
         """Count entities in this collection according to criteria.

src/aiida/orm/nodes/data/code/legacy.py

@@ -220,7 +220,7 @@ def get_code_helper(cls, label, machinename=None, backend=None):
         elif query.count() > 1:
             codes = query.all(flat=True)
             retstr = f"There are multiple codes with label '{label}', having IDs: "
-            retstr += f"{', '.join(sorted([str(c.pk) for c in codes]))}.\n"  # type: ignore[union-attr]
+            retstr += f"{', '.join(sorted([str(c.pk) for c in codes]))}.\n"
             retstr += 'Relabel them (using their ID), or refer to them with their ID.'
             raise MultipleObjectsError(retstr)
         else:
@@ -320,9 +320,9 @@ def list_for_plugin(cls, plugin, labels=True, backend=None):
         valid_codes = query.all(flat=True)
 
         if labels:
-            return [c.label for c in valid_codes]  # type: ignore[union-attr]
+            return [c.label for c in valid_codes]
 
-        return [c.pk for c in valid_codes]  # type: ignore[union-attr]
+        return [c.pk for c in valid_codes]
 
     def _validate(self):
         super()._validate()

src/aiida/orm/querybuilder.py

@@ -1033,12 +1033,12 @@ def _get_aiida_entity_res(value) -> Any:
             return value
 
     @overload
-    def first(self, flat: Literal[False] = False) -> Optional[list[Any]]: ...
+    def first(self, flat: Literal[False] = False) -> list[Any] | None: ...
 
     @overload
-    def first(self, flat: Literal[True]) -> Optional[Any]: ...
+    def first(self, flat: Literal[True]) -> Any | None: ...
 
-    def first(self, flat: bool = False) -> Optional[list[Any] | Any]:
+    def first(self, flat: bool = False) -> list[Any] | Any | None:
         """Return the first result of the query.
 
         Calling ``first`` results in an execution of the underlying query.
@@ -1105,7 +1105,13 @@ def iterdict(self, batch_size: Optional[int] = 100) -> Iterable[Dict[str, Dict[s
 
             yield item
 
-    def all(self, batch_size: Optional[int] = None, flat: bool = False) -> Union[List[List[Any]], List[Any]]:
+    @overload
+    def all(self, batch_size: int | None = None, flat: Literal[False] = False) -> list[list[Any]]: ...
+
+    @overload
+    def all(self, batch_size: int | None = None, flat: Literal[True] = True) -> list[Any]: ...
+
+    def all(self, batch_size: int | None = None, flat: bool = False) -> list[list[Any]] | list[Any]:
         """Executes the full query with the order of the rows as returned by the backend.
 
         The order inside each row is given by the order of the vertices in the path and the order of the projections for

src/aiida/tools/_dumping/detect.py

@@ -11,7 +11,7 @@
 from __future__ import annotations
 
 from datetime import datetime, timedelta
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Tuple, Type, Union, cast
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type, Union, cast
 
 from aiida import orm
 from aiida.common import AIIDA_LOGGER
@@ -171,7 +171,7 @@ def _query_single_type(
         if self.NODE_TAG not in qb._projections:
             qb.add_projection(self.NODE_TAG, '*')
 
-        return cast(list[orm.ProcessNode], qb.all(flat=True))
+        return qb.all(flat=True)
 
     def _exclude_tracked_nodes(self, nodes: list[orm.ProcessNode], store_type: str) -> list[orm.ProcessNode]:
         """Exclude nodes that are already tracked in the dump tracker.
@@ -289,7 +289,7 @@ def _detect_deleted_nodes(self) -> set[str]:
             qb = orm.QueryBuilder()
             orm_type = REGISTRY_TO_ORM_TYPE[registry_name]
             qb.append(orm_type, project=['uuid'])
-            all_db_uuids = cast(Set[str], set(qb.all(flat=True)))
+            all_db_uuids = set(qb.all(flat=True))
 
             # Find missing UUIDs
             missing_uuids = dumped_uuids - all_db_uuids

src/aiida/tools/_dumping/executors/profile.py

@@ -11,7 +11,7 @@
 from __future__ import annotations
 
 from pathlib import Path
-from typing import TYPE_CHECKING, List, cast
+from typing import TYPE_CHECKING, cast
 
 from aiida import orm
 from aiida.common import NotExistent
@@ -91,7 +91,7 @@ def _determine_groups_to_process(self) -> list[orm.Group]:
         """Determine which groups to process based on config."""
         if self.config.all_entries:
             qb_groups = orm.QueryBuilder().append(orm.Group)
-            return cast(List[orm.Group], qb_groups.all(flat=True))
+            return qb_groups.all(flat=True)
 
         if not self.config.groups:
             return []

src/aiida/tools/archive/create.py

@@ -16,7 +16,7 @@
 import tempfile
 from datetime import datetime
 from pathlib import Path
-from typing import Callable, Dict, Iterable, List, Optional, Sequence, Set, Tuple, Union
+from typing import Callable, Iterable, Optional, Sequence, Union
 
 from tabulate import tabulate
 
@@ -192,7 +192,7 @@ def querybuilder():
     EXPORT_LOGGER.report(initial_summary)
 
     # Store starting UUIDs, to write to metadata
-    starting_uuids: Dict[EntityTypes, Set[str]] = {
+    starting_uuids: dict[EntityTypes, set[str]] = {
         EntityTypes.USER: set(),
         EntityTypes.COMPUTER: set(),
         EntityTypes.GROUP: set(),
@@ -201,7 +201,7 @@ def querybuilder():
 
     # Store all entity IDs to be written to the archive
     # Note, this is the order they will be written to the archive
-    entity_ids: Dict[EntityTypes, Set[int]] = {
+    entity_ids: dict[EntityTypes, set[int]] = {
         ent: set()
         for ent in [
             EntityTypes.USER,
@@ -376,12 +376,12 @@ def transform(d):
 
 def _collect_all_entities(
     querybuilder: QbType,
-    entity_ids: Dict[EntityTypes, Set[int]],
+    entity_ids: dict[EntityTypes, set[int]],
     include_authinfos: bool,
     include_comments: bool,
     include_logs: bool,
     batch_size: int,
-) -> Tuple[List[Tuple[int, int]], Set[LinkQuadruple]]:
+) -> tuple[list[tuple[int, int]], set[LinkQuadruple]]:
     """Collect all entities.
 
     :returns: (group_id_to_node_id, link_data) and updates entity_ids
@@ -393,11 +393,7 @@ def progress_str(name):
     with get_progress_reporter()(desc=progress_str(''), total=9) as progress:
         progress.set_description_str(progress_str('Nodes'))
         entity_ids[EntityTypes.NODE].update(
-            querybuilder()
-            .append(orm.Node, project='id')
-            .all(  # type: ignore[arg-type]
-                batch_size=batch_size, flat=True
-            )
+            querybuilder().append(orm.Node, project='id').all(batch_size=batch_size, flat=True)
         )
         progress.update()
 
@@ -417,7 +413,7 @@ def progress_str(name):
             querybuilder()
             .append(
                 orm.Group,
-                project='id',  # type: ignore[arg-type]
+                project='id',
             )
             .all(batch_size=batch_size, flat=True)
         )
@@ -429,15 +425,15 @@ def progress_str(name):
             .append(orm.Node, with_group='group', project='id')
             .distinct()
         )
-        group_nodes: List[Tuple[int, int]] = qbuilder.all(batch_size=batch_size)  # type: ignore[assignment]
+        group_nodes: list[tuple[int, int]] = qbuilder.all(batch_size=batch_size)  # type: ignore[assignment]
 
         progress.set_description_str(progress_str('Computers'))
         progress.update()
         entity_ids[EntityTypes.COMPUTER].update(
             querybuilder()
             .append(
                 orm.Computer,
-                project='id',  # type: ignore[arg-type]
+                project='id',
             )
             .all(batch_size=batch_size, flat=True)
         )
@@ -449,7 +445,7 @@ def progress_str(name):
                 querybuilder()
                 .append(
                     orm.AuthInfo,
-                    project='id',  # type: ignore[arg-type]
+                    project='id',
                 )
                 .all(batch_size=batch_size, flat=True)
             )
@@ -461,7 +457,7 @@ def progress_str(name):
                 querybuilder()
                 .append(
                     orm.Log,
-                    project='id',  # type: ignore[arg-type]
+                    project='id',
                 )
                 .all(batch_size=batch_size, flat=True)
             )
@@ -473,7 +469,7 @@ def progress_str(name):
                 querybuilder()
                 .append(
                     orm.Comment,
-                    project='id',  # type: ignore[arg-type]
+                    project='id',
                 )
                 .all(batch_size=batch_size, flat=True)
             )
@@ -484,7 +480,7 @@ def progress_str(name):
             querybuilder()
             .append(
                 orm.User,
-                project='id',  # type: ignore[arg-type]
+                project='id',
             )
             .all(batch_size=batch_size, flat=True)
         )
@@ -494,14 +490,14 @@ def progress_str(name):
 
 def _collect_required_entities(
     querybuilder: QbType,
-    entity_ids: Dict[EntityTypes, Set[int]],
-    traversal_rules: Dict[str, bool],
+    entity_ids: dict[EntityTypes, set[int]],
+    traversal_rules: dict[str, bool],
     include_authinfos: bool,
     include_comments: bool,
     include_logs: bool,
     backend: StorageBackend,
     batch_size: int,
-) -> Tuple[List[Tuple[int, int]], Set[LinkQuadruple]]:
+) -> tuple[list[tuple[int, int]], set[LinkQuadruple]]:
     """Collect required entities, given a set of starting entities and provenance graph traversal rules.
 
     :returns: (group_id_to_node_id, link_data) and updates entity_ids
@@ -513,7 +509,7 @@ def progress_str(name):
     with get_progress_reporter()(desc=progress_str(''), total=7) as progress:
         # get all nodes from groups
         progress.set_description_str(progress_str('Nodes (groups)'))
-        group_nodes: List[Tuple[int, int]] = []
+        group_nodes: list[tuple[int, int]] = []
         if entity_ids[EntityTypes.GROUP]:
             qbuilder = querybuilder()
             qbuilder.append(
@@ -632,7 +628,7 @@ def progress_str(name):
 
 
 def _stream_repo_files(
-    key_format: str, writer: ArchiveWriterAbstract, node_ids: Set[int], backend: StorageBackend, batch_size: int
+    key_format: str, writer: ArchiveWriterAbstract, node_ids: set[int], backend: StorageBackend, batch_size: int
 ) -> None:
     """Collect all repository object keys from the nodes, then stream the files to the archive."""
     keys = set(
@@ -652,7 +648,7 @@ def _stream_repo_files(
             progress.update()
 
 
-def _check_unsealed_nodes(querybuilder: QbType, node_ids: Set[int], batch_size: int) -> None:
+def _check_unsealed_nodes(querybuilder: QbType, node_ids: set[int], batch_size: int) -> None:
     """Check no process nodes are unsealed, i.e. all processes have completed."""
     qbuilder = (
         querybuilder()
@@ -678,7 +674,7 @@ def _check_unsealed_nodes(querybuilder: QbType, node_ids: Set[int], batch_size:
 
 def _check_node_licenses(
     querybuilder: QbType,
-    node_ids: Set[int],
+    node_ids: set[int],
     allowed_licenses: Union[None, Sequence[str], Callable],
     forbidden_licenses: Union[None, Sequence[str], Callable],
     batch_size: int,