1from typing import (
2    Callable,
3    Dict,
4    Generic,
5    Iterable,
6    List,
7    Mapping,
8    Optional,
9    Tuple,
10    Type,
11    TypeVar,
12    Union,
13    Sequence,
14)
15
16K = TypeVar('K')
17T = TypeVar('T')
18T_co = TypeVar('T_co', covariant=True)
19V = TypeVar('V')
20
21
22just_float = 42.  # type: float
23optional_float = 42.  # type: Optional[float]
24list_of_ints = [42]  # type: List[int]
25list_of_floats = [42.]  # type: List[float]
26list_of_optional_floats = [x or None for x in list_of_floats]  # type: List[Optional[float]]
27list_of_ints_and_strs = [42, 'abc']  # type: List[Union[int, str]]
28
29# Test that simple parameters are handled
30def list_t_to_list_t(the_list: List[T]) -> List[T]:
31    return the_list
32
33x0 = list_t_to_list_t(list_of_ints)[0]
34#? int()
35x0
36
37for a in list_t_to_list_t(list_of_ints):
38    #? int()
39    a
40
41# Test that unions are handled
42x2 = list_t_to_list_t(list_of_ints_and_strs)[0]
43#? int() str()
44x2
45
46for z in list_t_to_list_t(list_of_ints_and_strs):
47    #? int() str()
48    z
49
50
51list_of_int_type = [int]  # type: List[Type[int]]
52
53# Test that nested parameters are handled
54def list_optional_t_to_list_t(the_list: List[Optional[T]]) -> List[T]:
55    return [x for x in the_list if x is not None]
56
57
58for xa in list_optional_t_to_list_t(list_of_optional_floats):
59    #? float()
60    xa
61
62# Under covariance rules this is strictly incorrect (because List is mutable,
63# the function would be allowed to put `None`s into our List[float], which would
64# be bad), however we don't expect jedi to enforce that.
65for xa1 in list_optional_t_to_list_t(list_of_floats):
66    #? float()
67    xa1
68
69
70def optional_t_to_list_t(x: Optional[T]) -> List[T]:
71    return [x] if x is not None else []
72
73
74for xb in optional_t_to_list_t(optional_float):
75    #? float()
76    xb
77
78
79for xb2 in optional_t_to_list_t(just_float):
80    #? float()
81    xb2
82
83
84def optional_list_t_to_list_t(x: Optional[List[T]]) -> List[T]:
85    return x if x is not None else []
86
87
88optional_list_float = None  # type: Optional[List[float]]
89for xc in optional_list_t_to_list_t(optional_list_float):
90    #? float()
91    xc
92
93for xc2 in optional_list_t_to_list_t(list_of_floats):
94    #? float()
95    xc2
96
97
98def list_type_t_to_list_t(the_list: List[Type[T]]) -> List[T]:
99    return [x() for x in the_list]
100
101
102x1 = list_type_t_to_list_t(list_of_int_type)[0]
103#? int()
104x1
105
106
107for b in list_type_t_to_list_t(list_of_int_type):
108    #? int()
109    b
110
111
112# Test construction of nested generic tuple return parameters
113def list_t_to_list_tuple_t(the_list: List[T]) -> List[Tuple[T]]:
114    return [(x,) for x in the_list]
115
116
117x1t = list_t_to_list_tuple_t(list_of_ints)[0][0]
118#? int()
119x1t
120
121
122for c1 in list_t_to_list_tuple_t(list_of_ints):
123    #? int()
124    c1[0]
125
126
127for c2, in list_t_to_list_tuple_t(list_of_ints):
128    #? int()
129    c2
130
131
132# Test handling of nested tuple input parameters
133def list_tuple_t_to_tuple_list_t(the_list: List[Tuple[T]]) -> Tuple[List[T], ...]:
134    return tuple(list(x) for x in the_list)
135
136
137list_of_int_tuples = [(x,) for x in list_of_ints]  # type: List[Tuple[int]]
138
139for b in list_tuple_t_to_tuple_list_t(list_of_int_tuples):
140    #? int()
141    b[0]
142
143
144def list_tuple_t_elipsis_to_tuple_list_t(the_list: List[Tuple[T, ...]]) -> Tuple[List[T], ...]:
145    return tuple(list(x) for x in the_list)
146
147
148list_of_int_tuple_elipsis = [tuple(list_of_ints)]  # type: List[Tuple[int, ...]]
149
150for b in list_tuple_t_elipsis_to_tuple_list_t(list_of_int_tuple_elipsis):
151    #? int()
152    b[0]
153
154
155# Test handling of nested callables
156def foo(x: int) -> int:
157    return x
158
159
160list_of_funcs = [foo]  # type: List[Callable[[int], int]]
161
162def list_func_t_to_list_func_type_t(the_list: List[Callable[[T], T]]) -> List[Callable[[Type[T]], T]]:
163    def adapt(func: Callable[[T], T]) -> Callable[[Type[T]], T]:
164        def wrapper(typ: Type[T]) -> T:
165            return func(typ())
166        return wrapper
167    return [adapt(x) for x in the_list]
168
169
170for b in list_func_t_to_list_func_type_t(list_of_funcs):
171    #? int()
172    b(int)
173
174
175def bar(*a, **k) -> int:
176    return len(a) + len(k)
177
178
179list_of_funcs_2 = [bar]  # type: List[Callable[..., int]]
180
181def list_func_t_passthrough(the_list: List[Callable[..., T]]) -> List[Callable[..., T]]:
182    return the_list
183
184
185for b in list_func_t_passthrough(list_of_funcs_2):
186    #? int()
187    b(None, x="x")
188
189
190mapping_int_str = {42: 'a'}  # type: Dict[int, str]
191
192# Test that mappings (that have more than one parameter) are handled
193def invert_mapping(mapping: Mapping[K, V]) -> Mapping[V, K]:
194    return {v: k for k, v in mapping.items()}
195
196#? int()
197invert_mapping(mapping_int_str)['a']
198
199
200# Test that the right type is chosen when a mapping is passed to something with
201# only a single parameter. This checks that our inheritance checking picks the
202# right thing.
203def first(iterable: Iterable[T]) -> T:
204    return next(iter(iterable))
205
206#? int()
207first(mapping_int_str)
208
209# Test inference of str as an iterable of str.
210#? str()
211first("abc")
212
213some_str = NotImplemented  # type: str
214#? str()
215first(some_str)
216
217annotated = [len]  # type: List[ Callable[[Sequence[float]], int] ]
218#? int()
219first(annotated)()
220
221# Test that the right type is chosen when a partially realised mapping is expected
222def values(mapping: Mapping[int, T]) -> List[T]:
223    return list(mapping.values())
224
225#? str()
226values(mapping_int_str)[0]
227
228x2 = values(mapping_int_str)[0]
229#? str()
230x2
231
232for b in values(mapping_int_str):
233    #? str()
234    b
235
236
237#
238# Tests that user-defined generic types are handled
239#
240list_ints = [42]  # type: List[int]
241
242class CustomGeneric(Generic[T_co]):
243    def __init__(self, val: T_co) -> None:
244        self.val = val
245
246
247# Test extraction of type from a custom generic type
248def custom(x: CustomGeneric[T]) -> T:
249    return x.val
250
251custom_instance = CustomGeneric(42)  # type: CustomGeneric[int]
252
253#? int()
254custom(custom_instance)
255
256x3 = custom(custom_instance)
257#? int()
258x3
259
260
261# Test construction of a custom generic type
262def wrap_custom(iterable: Iterable[T]) -> List[CustomGeneric[T]]:
263    return [CustomGeneric(x) for x in iterable]
264
265#? int()
266wrap_custom(list_ints)[0].val
267
268x4 = wrap_custom(list_ints)[0]
269#? int()
270x4.val
271
272for x5 in wrap_custom(list_ints):
273    #? int()
274    x5.val
275
276
277# Test extraction of type from a nested custom generic type
278list_custom_instances = [CustomGeneric(42)]  # type: List[CustomGeneric[int]]
279
280def unwrap_custom(iterable: Iterable[CustomGeneric[T]]) -> List[T]:
281    return [x.val for x in iterable]
282
283#? int()
284unwrap_custom(list_custom_instances)[0]
285
286x6 = unwrap_custom(list_custom_instances)[0]
287#? int()
288x6
289
290for x7 in unwrap_custom(list_custom_instances):
291    #? int()
292    x7
293
294
295for xc in unwrap_custom([CustomGeneric(s) for s in 'abc']):
296    #? str()
297    xc
298
299
300for xg in unwrap_custom(CustomGeneric(s) for s in 'abc'):
301    #? str()
302    xg
303
304
305# Test extraction of type from type parameer nested within a custom generic type
306custom_instance_list_int = CustomGeneric([42])  # type: CustomGeneric[List[int]]
307
308def unwrap_custom2(instance: CustomGeneric[Iterable[T]]) -> List[T]:
309    return list(instance.val)
310
311#? int()
312unwrap_custom2(custom_instance_list_int)[0]
313
314x8 = unwrap_custom2(custom_instance_list_int)[0]
315#? int()
316x8
317
318for x9 in unwrap_custom2(custom_instance_list_int):
319    #? int()
320    x9
321
322
323# Test that classes which have generic parents but are not generic themselves
324# are still inferred correctly.
325class Specialised(Mapping[int, str]):
326    pass
327
328
329specialised_instance = NotImplemented  # type: Specialised
330
331#? int()
332first(specialised_instance)
333
334#? str()
335values(specialised_instance)[0]
336
337
338# Test that classes which have generic ancestry but neither they nor their
339# parents are not generic are still inferred correctly.
340class ChildOfSpecialised(Specialised):
341    pass
342
343
344child_of_specialised_instance = NotImplemented  # type: ChildOfSpecialised
345
346#? int()
347first(child_of_specialised_instance)
348
349#? str()
350values(child_of_specialised_instance)[0]
351
352
353# Test that unbound generics are inferred as much as possible
354class CustomPartialGeneric1(Mapping[str, T]):
355    pass
356
357
358custom_partial1_instance = NotImplemented  # type: CustomPartialGeneric1[int]
359
360#? str()
361first(custom_partial1_instance)
362
363
364custom_partial1_unbound_instance = NotImplemented  # type: CustomPartialGeneric1
365
366#? str()
367first(custom_partial1_unbound_instance)
368
369
370class CustomPartialGeneric2(Mapping[T, str]):
371    pass
372
373
374custom_partial2_instance = NotImplemented  # type: CustomPartialGeneric2[int]
375
376#? int()
377first(custom_partial2_instance)
378
379#? str()
380values(custom_partial2_instance)[0]
381
382
383custom_partial2_unbound_instance = NotImplemented  # type: CustomPartialGeneric2
384
385#? []
386first(custom_partial2_unbound_instance)
387
388#? str()
389values(custom_partial2_unbound_instance)[0]
390