Define overloaded methods for typing purposes

aiida/orm/querybuilder.py

@@ -19,6 +19,8 @@
 An instance of one of the implementation classes becomes a member of the :func:`QueryBuilder` instance
 when instantiated by the user.
 """
+from __future__ import annotations
+
 from copy import deepcopy
 from inspect import isclass as inspect_isclass
 from typing import (
@@ -27,6 +29,7 @@
     Dict,
     Iterable,
     List,
+    Literal,
     NamedTuple,
     Optional,
     Sequence,
@@ -35,6 +38,7 @@
     Type,
     Union,
     cast,
+    overload,
 )
 import warnings
 
@@ -989,7 +993,15 @@ def _get_aiida_entity_res(value) -> Any:
         except TypeError:
             return value
 
-    def first(self, flat: bool = False) -> Optional[Union[List[Any], Any]]:
+    @overload
+    def first(self, flat: Literal[False]) -> Optional[list[Any]]:
+        ...
+
+    @overload
+    def first(self, flat: Literal[True]) -> Optional[Any]:
+        ...
+
+    def first(self, flat: bool = False) -> Optional[list[Any] | Any]:
         """Return the first result of the query.
 
         Calling ``first`` results in an execution of the underlying query.