1from __future__ import annotations
 2
 3from typing import Generic, TypeVar, cast
 4
 5from openai._utils import extract_type_var_from_base
 6
 7_T = TypeVar("_T")
 8_T2 = TypeVar("_T2")
 9_T3 = TypeVar("_T3")
10
11
12class BaseGeneric(Generic[_T]): ...
13
14
15class SubclassGeneric(BaseGeneric[_T]): ...
16
17
18class BaseGenericMultipleTypeArgs(Generic[_T, _T2, _T3]): ...
19
20
21class SubclassGenericMultipleTypeArgs(BaseGenericMultipleTypeArgs[_T, _T2, _T3]): ...
22
23
24class SubclassDifferentOrderGenericMultipleTypeArgs(BaseGenericMultipleTypeArgs[_T2, _T, _T3]): ...
25
26
27def test_extract_type_var() -> None:
28    assert (
29        extract_type_var_from_base(
30            BaseGeneric[int],
31            index=0,
32            generic_bases=cast("tuple[type, ...]", (BaseGeneric,)),
33        )
34        == int
35    )
36
37
38def test_extract_type_var_generic_subclass() -> None:
39    assert (
40        extract_type_var_from_base(
41            SubclassGeneric[int],
42            index=0,
43            generic_bases=cast("tuple[type, ...]", (BaseGeneric,)),
44        )
45        == int
46    )
47
48
49def test_extract_type_var_multiple() -> None:
50    typ = BaseGenericMultipleTypeArgs[int, str, None]
51
52    generic_bases = cast("tuple[type, ...]", (BaseGenericMultipleTypeArgs,))
53    assert extract_type_var_from_base(typ, index=0, generic_bases=generic_bases) == int
54    assert extract_type_var_from_base(typ, index=1, generic_bases=generic_bases) == str
55    assert extract_type_var_from_base(typ, index=2, generic_bases=generic_bases) == type(None)
56
57
58def test_extract_type_var_generic_subclass_multiple() -> None:
59    typ = SubclassGenericMultipleTypeArgs[int, str, None]
60
61    generic_bases = cast("tuple[type, ...]", (BaseGenericMultipleTypeArgs,))
62    assert extract_type_var_from_base(typ, index=0, generic_bases=generic_bases) == int
63    assert extract_type_var_from_base(typ, index=1, generic_bases=generic_bases) == str
64    assert extract_type_var_from_base(typ, index=2, generic_bases=generic_bases) == type(None)
65
66
67def test_extract_type_var_generic_subclass_different_ordering_multiple() -> None:
68    typ = SubclassDifferentOrderGenericMultipleTypeArgs[int, str, None]
69
70    generic_bases = cast("tuple[type, ...]", (BaseGenericMultipleTypeArgs,))
71    assert extract_type_var_from_base(typ, index=0, generic_bases=generic_bases) == int
72    assert extract_type_var_from_base(typ, index=1, generic_bases=generic_bases) == str
73    assert extract_type_var_from_base(typ, index=2, generic_bases=generic_bases) == type(None)