1from functools import wraps
2
3from . import filters
4from .asyncsupport import auto_aiter
5from .asyncsupport import auto_await
6
7
8async def auto_to_seq(value):
9    seq = []
10    if hasattr(value, "__aiter__"):
11        async for item in value:
12            seq.append(item)
13    else:
14        for item in value:
15            seq.append(item)
16    return seq
17
18
19async def async_select_or_reject(args, kwargs, modfunc, lookup_attr):
20    seq, func = filters.prepare_select_or_reject(args, kwargs, modfunc, lookup_attr)
21    if seq:
22        async for item in auto_aiter(seq):
23            if func(item):
24                yield item
25
26
27def dualfilter(normal_filter, async_filter):
28    wrap_evalctx = False
29    if getattr(normal_filter, "environmentfilter", False) is True:
30
31        def is_async(args):
32            return args[0].is_async
33
34        wrap_evalctx = False
35    else:
36        has_evalctxfilter = getattr(normal_filter, "evalcontextfilter", False) is True
37        has_ctxfilter = getattr(normal_filter, "contextfilter", False) is True
38        wrap_evalctx = not has_evalctxfilter and not has_ctxfilter
39
40        def is_async(args):
41            return args[0].environment.is_async
42
43    @wraps(normal_filter)
44    def wrapper(*args, **kwargs):
45        b = is_async(args)
46        if wrap_evalctx:
47            args = args[1:]
48        if b:
49            return async_filter(*args, **kwargs)
50        return normal_filter(*args, **kwargs)
51
52    if wrap_evalctx:
53        wrapper.evalcontextfilter = True
54
55    wrapper.asyncfiltervariant = True
56
57    return wrapper
58
59
60def asyncfiltervariant(original):
61    def decorator(f):
62        return dualfilter(original, f)
63
64    return decorator
65
66
67@asyncfiltervariant(filters.do_first)
68async def do_first(environment, seq):
69    try:
70        return await auto_aiter(seq).__anext__()
71    except StopAsyncIteration:
72        return environment.undefined("No first item, sequence was empty.")
73
74
75@asyncfiltervariant(filters.do_groupby)
76async def do_groupby(environment, value, attribute):
77    expr = filters.make_attrgetter(environment, attribute)
78    return [
79        filters._GroupTuple(key, await auto_to_seq(values))
80        for key, values in filters.groupby(
81            sorted(await auto_to_seq(value), key=expr), expr
82        )
83    ]
84
85
86@asyncfiltervariant(filters.do_join)
87async def do_join(eval_ctx, value, d=u"", attribute=None):
88    return filters.do_join(eval_ctx, await auto_to_seq(value), d, attribute)
89
90
91@asyncfiltervariant(filters.do_list)
92async def do_list(value):
93    return await auto_to_seq(value)
94
95
96@asyncfiltervariant(filters.do_reject)
97async def do_reject(*args, **kwargs):
98    return async_select_or_reject(args, kwargs, lambda x: not x, False)
99
100
101@asyncfiltervariant(filters.do_rejectattr)
102async def do_rejectattr(*args, **kwargs):
103    return async_select_or_reject(args, kwargs, lambda x: not x, True)
104
105
106@asyncfiltervariant(filters.do_select)
107async def do_select(*args, **kwargs):
108    return async_select_or_reject(args, kwargs, lambda x: x, False)
109
110
111@asyncfiltervariant(filters.do_selectattr)
112async def do_selectattr(*args, **kwargs):
113    return async_select_or_reject(args, kwargs, lambda x: x, True)
114
115
116@asyncfiltervariant(filters.do_map)
117async def do_map(*args, **kwargs):
118    seq, func = filters.prepare_map(args, kwargs)
119    if seq:
120        async for item in auto_aiter(seq):
121            yield await auto_await(func(item))
122
123
124@asyncfiltervariant(filters.do_sum)
125async def do_sum(environment, iterable, attribute=None, start=0):
126    rv = start
127    if attribute is not None:
128        func = filters.make_attrgetter(environment, attribute)
129    else:
130
131        def func(x):
132            return x
133
134    async for item in auto_aiter(iterable):
135        rv += func(item)
136    return rv
137
138
139@asyncfiltervariant(filters.do_slice)
140async def do_slice(value, slices, fill_with=None):
141    return filters.do_slice(await auto_to_seq(value), slices, fill_with)
142
143
144ASYNC_FILTERS = {
145    "first": do_first,
146    "groupby": do_groupby,
147    "join": do_join,
148    "list": do_list,
149    # we intentionally do not support do_last because that would be
150    # ridiculous
151    "reject": do_reject,
152    "rejectattr": do_rejectattr,
153    "map": do_map,
154    "select": do_select,
155    "selectattr": do_selectattr,
156    "sum": do_sum,
157    "slice": do_slice,
158}
159