From d0f88ea563672ef5b746c7dbdcce32c28a884b23 Mon Sep 17 00:00:00 2001
From: Gabriel Gerlero <gerlero@users.noreply.github.com>
Date: Wed, 20 Mar 2024 10:30:52 -0300
Subject: [PATCH] Update type hints

---
 foamlib/_dictionaries.py   | 25 +++++++++++++++++++------
 tests/test_dictionaries.py |  2 +-
 2 files changed, 20 insertions(+), 7 deletions(-)

diff --git a/foamlib/_dictionaries.py b/foamlib/_dictionaries.py
index b1e53df..1a0e9bc 100644
--- a/foamlib/_dictionaries.py
+++ b/foamlib/_dictionaries.py
@@ -15,11 +15,14 @@
 
 from ._subprocesses import run_process, CalledProcessError
 
-np: Optional[Any]
 try:
     import numpy as np
+    from numpy.typing import NDArray
 except ModuleNotFoundError:
-    np = None
+    numpy = False
+else:
+    numpy = True
+
 
 FoamDimensionSet = namedtuple(
     "FoamDimensionSet",
@@ -177,7 +180,7 @@ def _serialize_sequence(sequence: Any) -> str:
     if (
         isinstance(sequence, Sequence)
         and not isinstance(sequence, str)
-        or np
+        or numpy
         and isinstance(sequence, np.ndarray)
     ):
         return f"({' '.join(_serialize(v) for v in sequence)})"
@@ -216,7 +219,7 @@ def _serialize_dimensions(value: Any) -> str:
     if (
         isinstance(value, Sequence)
         and not isinstance(value, str)
-        or np
+        or numpy
         and isinstance(value, np.ndarray)
     ) and len(value) == 7:
         return f"[{' '.join(str(v) for v in value)}]"
@@ -372,7 +375,12 @@ def type(self, value: str) -> None:
     @property
     def value(
         self,
-    ) -> Union[int, float, Sequence[Union[int, float, Sequence[Union[int, float]]]]]:
+    ) -> Union[
+        int,
+        float,
+        Sequence[Union[int, float, Sequence[Union[int, float]]]],
+        NDArray[np.generic],
+    ]:
         """
         Alias of `self["value"]`.
         """
@@ -457,7 +465,12 @@ def dimensions(
     @property
     def internal_field(
         self,
-    ) -> Union[int, float, Sequence[Union[int, float, Sequence[Union[int, float]]]]]:
+    ) -> Union[
+        int,
+        float,
+        Sequence[Union[int, float, Sequence[Union[int, float]]]],
+        NDArray[np.generic],
+    ]:
         """
         Alias of `self["internalField"]`.
         """
diff --git a/tests/test_dictionaries.py b/tests/test_dictionaries.py
index e7d4dc5..e0a006f 100644
--- a/tests/test_dictionaries.py
+++ b/tests/test_dictionaries.py
@@ -140,7 +140,7 @@ def test_internal_field(pitz: FoamCase) -> None:
     p_arr = np.arange(size) * 1e-6
     U_arr = np.full((size, 3), [-1e-6, 1e-6, 0]) * np.arange(size)[:, np.newaxis]
 
-    pitz[0]["p"].internal_field = p_arr  # type: ignore
+    pitz[0]["p"].internal_field = p_arr
     pitz[0]["U"].internal_field = U_arr
 
     assert pitz[0]["p"].internal_field == pytest.approx(p_arr)