fix lint

Author: ConvolutedDog

nsight/extraction.py

@@ -24,6 +24,7 @@
 import ncu_report
 import numpy as np
 import pandas as pd
+from numpy.typing import NDArray
 
 from nsight import exceptions, utils
 from nsight.utils import is_scalar
@@ -105,7 +106,7 @@ def extract_df_from_report(
         )
 
     annotations: List[str] = []
-    all_values: List[np.ndarray | None] = []
+    all_values: List[NDArray[Any] | None] = []
     kernel_names: List[str] = []
     gpus: List[str] = []
     compute_clocks: List[int] = []

nsight/transformation.py

@@ -14,9 +14,10 @@
 
 import numpy as np
 import pandas as pd
+from numpy.typing import NDArray
 
 
-def _value_aggregator(agg_func_name: str) -> Callable[[pd.Series], np.ndarray]:
+def _value_aggregator(agg_func_name: str) -> Callable[[pd.Series], NDArray[Any]]:
     """Factory function to create value aggregators.
 
     Args:
@@ -44,7 +45,7 @@ def _value_aggregator(agg_func_name: str) -> Callable[[pd.Series], np.ndarray]:
 
     numpy_agg_func = AGG_FUNCTIONS[agg_func_name]
 
-    def aggregator(series: pd.Series) -> np.ndarray:
+    def aggregator(series: pd.Series) -> NDArray[Any]:
         # Convert None to np.nan
         cleaned_series = series.apply(lambda x: np.nan if x is None else x)
         # Convert to numpy array, handling tuples

nsight/utils.py

@@ -13,6 +13,7 @@
 from typing import Any, Iterator
 
 import numpy as np
+from numpy.typing import NDArray
 
 from nsight.exceptions import (
     CUDA_CORE_UNAVAILABLE_MSG,
@@ -133,7 +134,7 @@ def print_header(*lines: str) -> None:
 @dataclass
 class NCUActionData:
     name: str
-    values: np.ndarray | None
+    values: NDArray[Any] | None
     compute_clock: int
     memory_clock: int
     gpu: str