Skip to content

Commit

Permalink
Use numpy typing
Browse files Browse the repository at this point in the history
  • Loading branch information
cphyc committed Sep 6, 2023
1 parent 9384474 commit da24357
Showing 1 changed file with 4 additions and 3 deletions.
7 changes: 4 additions & 3 deletions yt/frontends/rockstar/data_structures.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
from typing import List, Optional

import numpy as np
from numpy import typing as npt

from yt.data_objects.static_output import ParticleDataset
from yt.frontends.halo_catalog.data_structures import HaloCatalogFile
Expand All @@ -20,7 +21,7 @@ class RockstarBinaryFile(HaloCatalogFile):
header: dict
_position_offset: int
_member_offset: int
_Npart: np.array
_Npart: npt.NDArray[np.int64]
_ids_halos: List[int]
_file_size: int

Expand All @@ -46,7 +47,7 @@ def __init__(self, ds, io, filename, file_id, range):

super().__init__(ds, io, filename, file_id, range)

def _read_member(self, ihalo: int) -> Optional[np.array]:
def _read_member(self, ihalo: int) -> Optional[npt.NDArray[np.int64]]:
if ihalo not in self._ids_halos:
return None

Expand All @@ -59,7 +60,7 @@ def _read_member(self, ihalo: int) -> Optional[np.array]:
ids = np.fromfile(f, dtype=np.int64, count=self._Npart[ind_halo])
return ids

def _read_particle_positions(self, ptype, f=None):
def _read_particle_positions(self, ptype: str, f=None):
"""
Read all particle positions in this file.
"""
Expand Down

0 comments on commit da24357

Please sign in to comment.