25
25
if TYPE_CHECKING :
26
26
from yt .frontends .ramses .data_structures import RAMSESDomainSubset
27
27
28
+
28
29
PARTICLE_HANDLERS : set [type ["ParticleFileHandler" ]] = set ()
29
30
30
31
@@ -65,20 +66,23 @@ class ParticleFileHandler(abc.ABC, HandlerMixin):
65
66
known_fields : list [FieldKey ]
66
67
67
68
# The function to employ to read the file
69
+ # NOTE: We omit the `ParticleFileHandler` argument since
70
+ # it is accessed as a method (so the first argument is
71
+ # assumed to be `self`).
68
72
reader : Callable [
69
- ["ParticleFileHandler" , " RAMSESDomainSubset" , list [tuple [ str , str ] ], int ],
70
- dict [tuple [ str , str ] , np .ndarray ],
73
+ ["RAMSESDomainSubset" , list [FieldKey ], int ],
74
+ dict [FieldKey , np .ndarray ],
71
75
]
72
76
73
77
# Name of the config section (if any)
74
78
config_field : str | None = None
75
79
76
80
## These properties are computed dynamically
77
81
# Mapping from field to offset in file
78
- _field_offsets : dict [tuple [ str , str ] , int ]
82
+ _field_offsets : dict [FieldKey , int ]
79
83
80
84
# Mapping from field to the type of the data (float, integer, ...)
81
- _field_types : dict [tuple [ str , str ] , str ]
85
+ _field_types : dict [FieldKey , str ]
82
86
83
87
# Number of particle in the domain
84
88
_local_particle_count : int
@@ -131,14 +135,14 @@ def read_header(self):
131
135
pass
132
136
133
137
@property
134
- def field_offsets (self ) -> dict [tuple [ str , str ] , int ]:
138
+ def field_offsets (self ) -> dict [FieldKey , int ]:
135
139
if hasattr (self , "_field_offsets" ):
136
140
return self ._field_offsets
137
141
self .read_header ()
138
142
return self ._field_offsets
139
143
140
144
@property
141
- def field_types (self ) -> dict [tuple [ str , str ] , str ]:
145
+ def field_types (self ) -> dict [FieldKey , str ]:
142
146
if hasattr (self , "_field_types" ):
143
147
return self ._field_types
144
148
self .read_header ()
@@ -158,9 +162,7 @@ def header(self) -> dict[str, Any]:
158
162
self .read_header ()
159
163
return self ._header
160
164
161
- def handle_field (
162
- self , field : tuple [str , str ], data_dict : dict [tuple [str , str ], np .ndarray ]
163
- ):
165
+ def handle_field (self , field : FieldKey , data_dict : dict [FieldKey , np .ndarray ]):
164
166
"""
165
167
This function allows custom code to be called to handle special cases,
166
168
such as the particle birth time.
@@ -169,9 +171,9 @@ def handle_field(
169
171
170
172
Parameters
171
173
----------
172
- field : tuple[str, str]
174
+ field : FieldKey
173
175
The field name.
174
- data_dict : dict[tuple[str, str] , np.ndarray]
176
+ data_dict : dict[FieldKey , np.ndarray]
175
177
A dictionary containing the data.
176
178
177
179
By default, this function does nothing.
@@ -344,9 +346,7 @@ def birth_file_fname(self):
344
346
def has_birth_file (self ):
345
347
return os .path .exists (self .birth_file_fname )
346
348
347
- def handle_field (
348
- self , field : tuple [str , str ], data_dict : dict [tuple [str , str ], np .ndarray ]
349
- ):
349
+ def handle_field (self , field : FieldKey , data_dict : dict [FieldKey , np .ndarray ]):
350
350
_ptype , fname = field
351
351
if not (fname == "particle_birth_time" and self .ds .cosmological_simulation ):
352
352
return
@@ -492,9 +492,7 @@ def read_header(self):
492
492
self ._field_offsets = field_offsets
493
493
self ._field_types = _pfields
494
494
495
- def handle_field (
496
- self , field : tuple [str , str ], data_dict : dict [tuple [str , str ], np .ndarray ]
497
- ):
495
+ def handle_field (self , field : FieldKey , data_dict : dict [FieldKey , np .ndarray ]):
498
496
_ptype , fname = field
499
497
if not (fname == "particle_birth_time" and self .ds .cosmological_simulation ):
500
498
return
0 commit comments