typ: fix type annotations for poly roots methods

Author: oscarbenjamin

src/flint/flint_base/flint_base.pyi

@@ -80,7 +80,7 @@ class flint_poly(flint_elem, Generic[Telem]):
     def is_one(self) -> bool: ...
     def is_constant(self) -> bool: ...
     def is_gen(self) -> bool: ...
-    def roots(self) -> list[tuple[Telem, int]]: ...
+    def roots(self) -> list[tuple[Telem, int]] | list[Telem]: ...
     # Should be list[arb]:
     def real_roots(self) -> list[Any]: ...
     # Should be list[acb]:

src/flint/flint_base/flint_context.pyi

@@ -48,6 +48,3 @@ class PrecisionManager:
     def __call__(self, func: Callable[_P, _R]) -> Callable[_P, _R]: ...
     def __enter__(self) -> None: ...
     def __exit__(self, exc_type: object, value: object, traceback: object) -> None: ...
-
-
-thectx: FlintContext

src/flint/types/fmpq_poly.pyi

@@ -1,7 +1,8 @@
-from typing import overload, Any, Sequence
+from typing import overload, Sequence
 from flint.flint_base.flint_base import flint_poly
 from flint.types.fmpz import fmpz, ifmpz
 from flint.types.fmpq import fmpq, ifmpq
+from flint.types.acb import acb
 from flint.types.fmpz_poly import fmpz_poly, ifmpz_poly
 
 
@@ -81,12 +82,12 @@ class fmpq_poly(flint_poly[fmpq]):
 
     def factor(self, *, monic: bool = False) -> tuple[fmpq, list[tuple[fmpq_poly, int]]]: ...
     def factor_squarefree(self) -> tuple[fmpq, list[tuple[fmpq_poly, int]]]: ...
+    def roots(self) -> list[tuple[fmpq, int]]: ...
 
     def sqrt(self) -> fmpq_poly: ...
 
     def deflation(self) -> tuple[fmpq_poly, int]: ...
-    # XXX: acb
-    def complex_roots(self, **kwargs) -> list[Any]: ...
+    def complex_roots(self, **kwargs) -> list[acb]: ...
 
     @staticmethod
     def bernoulli_poly(n: int, /) -> fmpq_poly: ...

src/flint/types/fmpz_poly.pyi

@@ -1,4 +1,4 @@
-from typing import overload, Any, Sequence
+from typing import overload, Sequence
 from flint.flint_base.flint_base import flint_poly
 from flint.types.fmpz import fmpz, ifmpz
 from flint.types.fmpq import fmpq
@@ -71,15 +71,15 @@ class fmpz_poly(flint_poly[fmpz]):
     def resultant(self, other: ifmpz_poly, /) -> fmpz: ...
     def factor(self) -> tuple[fmpz, list[tuple[fmpz_poly, int]]]: ...
     def factor_squarefree(self) -> tuple[fmpz, list[tuple[fmpz_poly, int]]]: ...
+    def roots(self) -> list[tuple[fmpz, int]]: ...
     def sqrt(self) -> fmpz_poly: ...
     def deflation(self) -> tuple[fmpz_poly, int]: ...
     def deflation_monom(self) -> tuple[fmpz_poly, int, fmpz_poly]: ...
     def inflate(self, n: int, /) -> fmpz_poly: ...
     def deflate(self, n: int, /) -> fmpz_poly: ...
     def deflation_index(self) -> tuple[int, int]: ...
 
-    # XXX acb
-    def complex_roots(self, verbose: bool = False) -> list[Any]: ...
+    def complex_roots(self, verbose: bool = False) -> list[acb]: ...
     @staticmethod
     def cyclotomic(n: int, /) -> fmpz_poly: ...
     def is_cyclotomic(self) -> int: ...

src/flint/typing.py

@@ -130,7 +130,7 @@ def is_zero(self) -> bool: ...
     def is_one(self) -> bool: ...
     def is_constant(self) -> bool: ...
     def is_gen(self) -> bool: ...
-    def roots(self) -> list[tuple[_Tscalar, int]]: ...
+    def roots(self) -> list[tuple[_Tscalar, int]] | list[_Tscalar]: ...
     # Should be list[arb]:
     def real_roots(self) -> list[Any]: ...
     # Should be list[acb]:
@@ -144,6 +144,7 @@ def sqrt(self) -> Self: ...
     def gcd(self, other: Self | _Tscalar, /) -> Self: ...
     def factor(self) -> tuple[_Tscalar, list[tuple[Self, int]]]: ...
     def factor_squarefree(self) -> tuple[_Tscalar, list[tuple[Self, int]]]: ...
+    def roots(self) -> list[tuple[_Tscalar, int]]: ...
     def deflation(self) -> tuple[Self, int]: ...
     def mul_low(self, other: Self, n: int) -> Self: ...
     def pow_trunc(self, e: int, n: int) -> Self: ...