11from __future__ import annotations
22
3- from typing import TYPE_CHECKING
3+ from typing import TYPE_CHECKING , cast
44
55import numpy as np
66import numpy .typing as npt
@@ -205,7 +205,7 @@ def eigen(
205205
206206 n = matrix .shape [0 ]
207207
208- scores = np .ones (n ) / n
208+ scores = np .ones (n , dtype = np . float64 ) / n
209209 scores_new = scores .copy ()
210210
211211 converged , iterations = False , 0
@@ -221,15 +221,15 @@ def eigen(
221221
222222 scores [:] = scores_new
223223
224- return scores , iterations
224+ return cast ( "npt.NDArray[np.float64]" , scores ) , iterations
225225
226226
227227def pagerank_matrix (
228- matrix : npt .NDArray [np .float64 ],
228+ matrix : npt .NDArray [np .floating [ Any ] ],
229229 damping : float ,
230- ) -> npt .NDArray [np .float64 ]:
230+ ) -> npt .NDArray [np .floating [ Any ] ]:
231231 if not matrix .size :
232- return np .zeros (0 , dtype = np . float64 )
232+ return np .zeros (0 , dtype = matrix . dtype )
233233
234234 p = 1. / int (matrix .shape [0 ])
235235
@@ -246,7 +246,7 @@ def pagerank(
246246 tolerance : float ,
247247 limit : int ,
248248) -> tuple [npt .NDArray [np .float64 ], int ]:
249- matrix = pagerank_matrix (matrix , damping )
249+ matrix = cast ( "npt.NDArray[np.float64]" , pagerank_matrix (matrix , damping ) )
250250
251251 scores , iterations = eigen (matrix , tolerance = tolerance , limit = limit )
252252 scores /= np .linalg .norm (scores , ord = 1 )
0 commit comments