From f3f7878521a271e66d3609167806dfb22885f0e7 Mon Sep 17 00:00:00 2001
From: Corentin Cadiou <corentin.cadiou@cphyc.me>
Date: Wed, 6 Sep 2023 10:22:27 +0100
Subject: [PATCH] Use numpy typing

---
 yt/frontends/rockstar/data_structures.py | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/yt/frontends/rockstar/data_structures.py b/yt/frontends/rockstar/data_structures.py
index 8baa377694..fe10460c84 100644
--- a/yt/frontends/rockstar/data_structures.py
+++ b/yt/frontends/rockstar/data_structures.py
@@ -1,7 +1,7 @@
 import glob
 import os
 from functools import cached_property
-from typing import List, Optional
+from typing import Any, List, Optional
 
 import numpy as np
 
@@ -20,7 +20,7 @@ class RockstarBinaryFile(HaloCatalogFile):
     header: dict
     _position_offset: int
     _member_offset: int
-    _Npart: np.array
+    _Npart: "np.ndarray[Any, np.dtype[np.int64]]"
     _ids_halos: List[int]
     _file_size: int
 
@@ -46,7 +46,9 @@ 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["np.ndarray[Any, np.dtype[np.int64]]"]:
         if ihalo not in self._ids_halos:
             return None
 
@@ -59,7 +61,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.
         """