Merge pull request #227 from jorenham/optimize.quadratic_assignment

✨ complete `optimize.quadratic_assignment`

Author: jorenham

scipy-stubs/optimize/_qap.pyi

@@ -1,7 +1,46 @@
-from typing import Final
+from typing import Final, Literal, TypedDict, overload, type_check_only
 
-from scipy._typing import Untyped
+import numpy as np
+import optype.numpy as onp
+from ._optimize import OptimizeResult as _OptimizeResult
+
+@type_check_only
+class _CommonOptions(TypedDict, total=False):
+    maximize: onp.ToBool
+    rng: onp.ToJustInt | np.random.Generator | np.random.RandomState | None
+    partial_match: onp.ToInt2D | None
+
+@type_check_only
+class _FAQOptions(_CommonOptions, TypedDict, total=False):
+    P0: onp.ToFloat2D | Literal["barycenter", "randomized"]
+    shuffle: onp.ToBool
+    maxiter: onp.ToJustInt
+    tol: onp.ToFloat
+
+@type_check_only
+class _2OptOptions(_CommonOptions, TypedDict, total=False):
+    partial_guess: onp.ToInt2D | None
+
+###
 
 QUADRATIC_ASSIGNMENT_METHODS: Final = ["faq", "2opt"]
 
-def quadratic_assignment(A: Untyped, B: Untyped, method: str = "faq", options: Untyped | None = None) -> Untyped: ...
+class OptimizeResult(_OptimizeResult):
+    col_ind: onp.Array1D[np.intp]
+    fun: float | np.float64
+    nit: int
+
+@overload
+def quadratic_assignment(
+    A: onp.ToFloat2D,
+    B: onp.ToFloat2D,
+    method: Literal["faq"] = "faq",
+    options: _FAQOptions | None = None,
+) -> OptimizeResult: ...
+@overload
+def quadratic_assignment(
+    A: onp.ToFloat2D,
+    B: onp.ToFloat2D,
+    method: Literal["2opt"],
+    options: _2OptOptions | None = None,
+) -> OptimizeResult: ...