Merge pull request #672 from gerlero/typing

Improve type hints

Author: gerlero

src/foamlib/_files/_parsing/parsing.py

@@ -1,7 +1,12 @@
 import dataclasses
 import sys
 from collections.abc import Collection, Iterator, Sequence
-from typing import cast
+from typing import cast, overload
+
+if sys.version_info >= (3, 11):
+    from typing import Never, Unpack
+else:
+    from typing_extensions import Never, Unpack
 
 if sys.version_info >= (3, 12):
     from typing import override
@@ -116,20 +121,63 @@ def _flatten_results(
                     )
         return ret
 
+    @overload
+    @with_default
+    def getall(self, keywords: tuple[()], /) -> Collection[StandaloneData]: ...
+
+    @overload
+    @with_default
+    def getall(
+        self, keywords: tuple[str, Unpack[tuple[str, ...]]], /
+    ) -> Collection[Data | EllipsisType | None]: ...
+
     @override
     @with_default
     def getall(
         self, keywords: tuple[str, ...], /
-    ) -> Collection[Data | StandaloneData | EllipsisType]:
+    ) -> Collection[Data | StandaloneData | EllipsisType | None]:
         return [entry.data for entry in self._parsed.getall(keywords)]
 
+    @overload
+    def __getitem__(self, keywords: tuple[()]) -> StandaloneData: ...
+
+    @overload
+    def __getitem__(
+        self, keywords: tuple[str, Unpack[tuple[str, ...]]]
+    ) -> Data | EllipsisType | None: ...
+
+    @override
+    def __getitem__(
+        self, keywords: tuple[str, ...]
+    ) -> Data | StandaloneData | EllipsisType | None:  # ty: ignore[invalid-method-override]
+        entry = self._parsed[keywords]
+        return entry.data
+
     @override
-    def __setitem__(
-        self, key: tuple[str, ...], value: Data | StandaloneData | EllipsisType | None
+    def __setitem__(  # ty: ignore[invalid-method-override]
+        self, key: Never, value: Never
     ) -> None:  # pragma: no cover
         msg = "Use 'put' method instead"
         raise NotImplementedError(msg)
 
+    @overload
+    def put(
+        self,
+        keywords: tuple[()],
+        /,
+        data: StandaloneData,
+        content: bytes,
+    ) -> None: ...
+
+    @overload
+    def put(
+        self,
+        keywords: tuple[str, Unpack[tuple[str, ...]]],
+        /,
+        data: Data | EllipsisType | None,
+        content: bytes,
+    ) -> None: ...
+
     def put(
         self,
         keywords: tuple[str, ...],
@@ -143,6 +191,24 @@ def put(
         self._parsed[keywords] = ParsedFile._Entry(data, start, start + len(content))
         self._remove_child_entries(keywords)
 
+    @overload
+    def add(
+        self,
+        keywords: tuple[()],
+        data: StandaloneData,
+        content: bytes,
+        /,
+    ) -> None: ...
+
+    @overload
+    def add(
+        self,
+        keywords: tuple[str, Unpack[tuple[str, ...]]],
+        data: Data | EllipsisType | None,
+        content: bytes,
+        /,
+    ) -> None: ...
+
     @override
     def add(  # ty: ignore[invalid-method-override]
         self,

src/foamlib/_files/_typing.py

@@ -1,6 +1,12 @@
+import sys
 from collections.abc import Mapping, Sequence
 from typing import Literal, TypeAlias
 
+if sys.version_info >= (3, 11):
+    from typing import Unpack
+else:
+    from typing_extensions import Unpack
+
 import numpy as np
 from multicollections import MultiDict
 
@@ -38,8 +44,11 @@
     DataEntry | Sequence["DataEntryLike | KeywordEntryLike | DictLike"] | FieldLike
 )
 
-Data: TypeAlias = DataEntry | tuple[DataEntry, ...]
-DataLike: TypeAlias = DataEntryLike | tuple[DataEntryLike, ...]
+Data: TypeAlias = DataEntry | tuple[DataEntry, DataEntry, Unpack[tuple[DataEntry, ...]]]
+DataLike: TypeAlias = (
+    DataEntryLike
+    | tuple[DataEntryLike, DataEntryLike, Unpack[tuple[DataEntryLike, ...]]]
+)
 
 StandaloneDataEntry: TypeAlias = (
     DataEntry
@@ -57,9 +66,21 @@
     | Sequence[Sequence[int]]
 )
 
-StandaloneData: TypeAlias = StandaloneDataEntry | tuple[StandaloneDataEntry, ...]
+StandaloneData: TypeAlias = (
+    StandaloneDataEntry
+    | tuple[
+        StandaloneDataEntry,
+        StandaloneDataEntry,
+        Unpack[tuple[StandaloneDataEntry, ...]],
+    ]
+)
 StandaloneDataLike: TypeAlias = (
-    StandaloneDataEntryLike | tuple[StandaloneDataEntryLike, ...]
+    StandaloneDataEntryLike
+    | tuple[
+        StandaloneDataEntryLike,
+        StandaloneDataEntryLike,
+        Unpack[tuple[StandaloneDataEntryLike, ...]],
+    ]
 )
 
 SubDict: TypeAlias = (