Fix mypy errors

Author: mx-moth

src/emsarray/conventions/ugrid.py

@@ -1006,7 +1006,7 @@ def edge_count(self) -> int:
             return self.dataset.sizes[self.edge_dimension]
 
         # By computing the edge_node array we can determine how many edges exist
-        return self.edge_node_array.shape[0]
+        return cast(int, self.edge_node_array.shape[0])
 
     @property
     def face_count(self) -> int:

src/emsarray/operations/cache.py

@@ -19,6 +19,7 @@
 """
 import hashlib
 import marshal
+from typing import cast
 
 import numpy
 import xarray
@@ -146,7 +147,7 @@ def make_cache_key(dataset: xarray.Dataset, hash: "hashlib._Hash | None" = None)
     and should not be relied upon.
     """
     if hash is None:
-        hash = hashlib.blake2b(digest_size=32)
+        hash = cast("hashlib._Hash", hashlib.blake2b(digest_size=32))
 
     dataset.ems.hash_geometry(hash)
 

src/emsarray/operations/geometry.py

@@ -7,7 +7,7 @@
 import pathlib
 from collections.abc import Generator, Iterable, Iterator
 from contextlib import contextmanager
-from typing import IO, Any, TypeVar
+from typing import IO, Any, Generic, TypeVar
 
 import geojson
 import shapefile
@@ -19,7 +19,7 @@
 T = TypeVar('T')
 
 
-class _dumpable_iterator(list):
+class _dumpable_iterator(Generic[T], list):
     """
     Wrap an iterator / generator so it can be used in `json.dumps()`.
     No guarantees that it works for anything else!

src/emsarray/operations/triangulate.py

@@ -158,15 +158,15 @@ def _add_triangles(face_index: int, vertex_triangles: numpy.ndarray) -> None:
         vertex_triangles = _triangulate_polygons_by_length(same_length_polygons)
 
         for face_index, triangles in zip(same_length_face_indices, vertex_triangles):
-            _add_triangles(face_index, triangles)
+            _add_triangles(int(face_index), triangles)
 
     # Triangulate each concave polygon using a slower manual method.
     # Anecdotally concave polygons are very rare,
     # so using a slower method isn't an issue.
     for face_index in polygon_is_concave:
         polygon = polygons[face_index]
         triangles = _triangulate_concave_polygon(polygon)
-        _add_triangles(face_index, triangles)
+        _add_triangles(int(face_index), triangles)
 
     # Check that we have handled each triangle we expected.
     assert current_face == total_triangles

src/emsarray/utils.py

@@ -228,7 +228,7 @@ def fix_time_units_for_ems(
     """
 
     with netCDF4.Dataset(dataset_path, 'r+') as dataset:
-        variable = dataset.variables[variable_name]
+        variable = dataset.variables[str(variable_name)]
 
         units = cast(str, variable.getncattr('units'))
         calendar = cast(str, variable.getncattr('calendar') or DEFAULT_CALENDAR)
@@ -668,7 +668,11 @@ def wind_dimension(
     return xarray.DataArray(data=new_data, dims=new_dims)
 
 
-def datetime_from_np_time(np_time: numpy.datetime64) -> datetime.datetime:
+def datetime_from_np_time(
+    np_time: numpy.datetime64,
+    *,
+    tz: datetime.tzinfo = datetime.timezone.utc,
+) -> datetime.datetime:
     """
     Convert a numpy :class:`~numpy.datetime64`
     to a python :class:`~datetime.datetime`.
@@ -682,8 +686,24 @@ def datetime_from_np_time(np_time: numpy.datetime64) -> datetime.datetime:
     A conversion that truncates data is not reported as an error.
     If you're using numpy datetime64 with attosecond accuracy,
     the Python datetime formatting methods are insufficient for your needs anyway.
+
+    Parameters
+    ==========
+    np_time : numpy.datetime64
+        The numpy datetime64 to convert to a Python datetime.
+    tz : datetime.tzinfo
+        The timezone that the numpy datetime is in.
+        Defaults to UTC, as xarray will convert all time variables to UTC when
+        opening files.
+
+    Returns
+    =======
+    datetime.datetime
+        A timezone aware Python datetime.datetime instance.
     """
-    return datetime.datetime.fromtimestamp(np_time.item() / 10**9)
+    epoc = numpy.datetime64('1970-01-01')
+    timestamp = (np_time - epoc).astype('timedelta64[ns]')
+    return datetime.datetime.fromtimestamp(timestamp.astype(float) / 1e9, tz=tz)
 
 
 class RequiresExtraException(Exception):