TYP: more type hints fixes for gadget

yt-project · Oct 8, 2021 · 7c974c7 · 7c974c7
1 parent 5b9dbe1
commit 7c974c7
Show file tree

Hide file tree

Showing 4 changed files with 20 additions and 18 deletions.
diff --git a/yt/frontends/gadget/data_structures.py b/yt/frontends/gadget/data_structures.py
@@ -212,7 +212,7 @@ def _initialize_frontend_specific(self):
 
 class GadgetDataset(SPHDataset):
     _index_class: Type[Index] = GadgetBinaryIndex
-    _file_class = GadgetBinaryFile
+    _file_class: Type = GadgetBinaryFile
     _field_info_class: Type[FieldInfoContainer] = GadgetFieldInfo
     _particle_mass_name = "Mass"
     _particle_coordinates_name = "Coordinates"

diff --git a/yt/frontends/gadget/io.py b/yt/frontends/gadget/io.py
@@ -292,7 +292,7 @@ def _identify_fields(self, data_file):
                     if not hasattr(g[kk], "shape"):
                         continue
                     if len(g[kk].shape) > 1:
-                        self._vector_fields[kk] = g[kk].shape[1]
+                        self._vector_fields_dict[kk] = g[kk].shape[1]
                     fields.append((ptype, str(kk)))
 
         f.close()
@@ -306,16 +306,22 @@ def _identify_fields(self, data_file):
 ZeroMass = object()
 
 
-class IOHandlerGadgetBinary(IOHandlerSPH):
-    _dataset_type = "gadget_binary"
-    _vector_fields = (
+_VECTOR_FIELDS = dict(
+    (
         ("Coordinates", 3),
         ("Velocity", 3),
         ("Velocities", 3),
         ("MagneticField", 3),
         ("FourMetalFractions", 4),
         ("ElevenMetalMasses", 11),
     )
+)
+
+
+class IOHandlerGadgetBinary(IOHandlerSPH):
+    _dataset_type = "gadget_binary"
+    _vector_fields = tuple(_VECTOR_FIELDS.keys())
+    _vector_fields_dict = _VECTOR_FIELDS
 
     # Particle types (Table 3 in GADGET-2 user guide)
     #
@@ -337,7 +343,6 @@ class IOHandlerGadgetBinary(IOHandlerSPH):
     _format = None
 
     def __init__(self, ds, *args, **kwargs):
-        self._vector_fields = dict(self._vector_fields)
         self._fields = ds._field_spec
         self._ptypes = ds._ptype_spec
         self.data_files = set()
@@ -437,15 +442,14 @@ def _read_field_from_file(self, f, count, name):
         else:
             dt = self._endian + self._float_type
         dt = np.dtype(dt)
-        if name in self._vector_fields:
-            count *= self._vector_fields[name]
+        count *= self._vector_fields_dict.get(name, 1)
         arr = np.fromfile(f, dtype=dt, count=count)
         # ensure data are in native endianness to avoid errors
         # when field data are passed to cython
         dt = dt.newbyteorder("N")
         arr = arr.astype(dt)
-        if name in self._vector_fields:
-            factor = self._vector_fields[name]
+        if name in self._vector_fields_dict:
+            factor = self._vector_fields_dict[name]
             arr = arr.reshape((count // factor, factor), order="C")
         return arr
 
@@ -530,14 +534,12 @@ def _calculate_field_offsets(
                 if (ptype, field) not in field_list:
                     continue
                 start_offset = df_start * fs
-                if field in self._vector_fields:
-                    start_offset *= self._vector_fields[field]
+                start_offset *= self._vector_fields_dict.get(field, 1)
                 pos += start_offset
                 offsets[(ptype, field)] = pos
                 any_ptypes = True
                 remain_offset = (pcount[ptype] - df_start) * fs
-                if field in self._vector_fields:
-                    remain_offset *= self._vector_fields[field]
+                remain_offset *= self._vector_fields_dict.get(field, 1)
                 pos += remain_offset
             pos += 4
             if not any_ptypes:

diff --git a/yt/frontends/gadget/testing.py b/yt/frontends/gadget/testing.py
@@ -4,7 +4,7 @@
 from .definitions import gadget_field_specs, gadget_ptype_specs
 from .io import IOHandlerGadgetBinary
 
-vector_fields = dict(IOHandlerGadgetBinary._vector_fields)
+vector_fields = dict(IOHandlerGadgetBinary._vector_fields_dict)
 
 block_ids = {
     "Coordinates": "POS",

diff --git a/yt/frontends/gadget_fof/fields.py b/yt/frontends/gadget_fof/fields.py
@@ -1,17 +1,17 @@
-from yt.fields.field_info_container import FieldInfoContainer
+from yt.fields.field_info_container import FieldInfoContainer, KnownFieldsT
 
 m_units = "code_mass"
 p_units = "code_length"
 v_units = "code_velocity"
 
 _pnums = 6
-_type_fields = tuple(
+_type_fields: KnownFieldsT = tuple(
     ("%s%sType_%d" % (ptype, field, pnum), (units, [], None))
     for pnum in range(_pnums)
     for field, units in (("Mass", m_units), ("Len", p_units))
     for ptype in ("Group", "Subhalo")
 )
-_sub_type_fields = tuple(
+_sub_type_fields: KnownFieldsT = tuple(
     ("Subhalo%sType_%d" % (field, pnum), (units, [], None))
     for pnum in range(_pnums)
     for field, units in (