async def merge(aiters):
aiters = aioitertools.iter(aiters)
iters = [None]
nexts = [asyncio.create_task(aioitertools.next(aiters))]
while len(iters) > 0:
await asyncio.wait(nexts, return_when=asyncio.FIRST_COMPLETED)
new_iters = []
completed_iters = set()
for i, future in enumerate(nexts):
if future.done():
try:
if iters[i] is None:
# new iterator
new_iters.append(aioitertools.iter(future.result()))
else:
# new item
yield future.result()
nexts[i] = asyncio.create_task(
aioitertools.next(
aiters if iters[i] is None else iters[i]))
except StopAsyncIteration:
completed_iters.add(i)
except concurrent.futures.CancelledError as e:
raise concurrent.futures.CancelledError(
'Future {!r} at index {!r} was cancelled'.format(
future, i)) from e
for i in sorted(completed_iters, reverse=True):
del iters[i]
del nexts[i]
iters += new_iters
nexts += [
asyncio.create_task(aioitertools.next(new_iter))
for new_iter in new_iters
]
async def test_iter_range(self):
it = ait.iter(srange)
self.assertIsInstance(it, AsyncIterator)
idx = 0
async for item in it:
self.assertEqual(item, srange[idx])
idx += 1
async def test_iter_async_generator(self):
async def async_gen():
yield 1
yield 2
agen = async_gen()
self.assertEqual(ait.iter(agen), agen)
async def unique_justseen(aiter, key=lambda x: x):
prev_result = object()
async for item in aioitertools.iter(aiter):
curr_result = key(item)
if curr_result != prev_result:
prev_result = curr_result
yield item
async def test_next_range(self):
it = ait.iter(srange)
self.assertEqual(await ait.next(it), 0)
self.assertEqual(await ait.next(it), 1)
self.assertEqual(await ait.next(it), 2)
with self.assertRaises(StopAsyncIteration):
await ait.next(it)
async def test_next_iterable(self):
class async_iter:
def __init__(self):
self.index = 0
def __aiter__(self):
return self
def __anext__(self):
if self.index > 2:
raise StopAsyncIteration()
return self.fake_next()
async def fake_next(self):
value = slist[self.index]
self.index += 1
return value
it = ait.iter(async_iter())
self.assertEqual(await ait.next(it), "A")
self.assertEqual(await ait.next(it), "B")
self.assertEqual(await ait.next(it), "C")
with self.assertRaises(StopAsyncIteration):
await ait.next(it)
it = iter(slist)
self.assertEqual(await ait.next(it), "A")
self.assertEqual(await ait.next(it), "B")
self.assertEqual(await ait.next(it), "C")
with self.assertRaises(StopAsyncIteration):
await ait.next(it)
async def test_iter_list(self):
it = ait.iter(slist)
self.assertIsInstance(it, AsyncIterator)
idx = 0
async for item in it:
self.assertEqual(item, slist[idx])
idx += 1
async def test_next_list(self):
it = ait.iter(slist)
self.assertEqual(await ait.next(it), "A")
self.assertEqual(await ait.next(it), "B")
self.assertEqual(await ait.next(it), "C")
with self.assertRaises(StopAsyncIteration):
await ait.next(it)
async def _func(it):
it = aioitertools.iter(it)
async for item in it:
if item:
yield {(query.new if k == query.old else k): v
for k, v in item.items()}
else:
yield item
async def _key_func(it):
it = aioitertools.iter(it)
heap = [(query.key(materialized), materialized)
async for materialized in it]
heapq.heapify(heap)
while heap:
_, materialized = heapq.heappop(heap)
yield materialized
async def union_iters(x):
result = {}
async for i in aioitertools.iter(x):
if isinstance(i, dict):
i = await materialize_walk(i)
result = union_dicts(result, i)
else:
raise Exception("non dict union is not yet supported")
return result
async def test_iter_iterable(self):
sentinel = object()
class async_iterable:
def __aiter__(self):
return sentinel
aiter = async_iterable()
self.assertEqual(ait.iter(aiter), sentinel)
async def _async_key_func(it):
it = aioitertools.iter(it)
materialized = [i async for i in it] # TODO: eliminate copy
keys = (query.key(item) for item in materialized)
keys = await asyncio.gather(*keys)
heap = [(y, x) for x, y in enumerate(keys)]
heapq.heapify(heap)
while heap:
_, index = heapq.heappop(heap)
yield materialized[index]
async def consume(aiter, n=float('inf')):
aiter = aioitertools.iter(aiter)
if n == float('inf'):
async for item in aiter:
pass
else:
for i in range(n):
try:
await aioitertools.next(aiter)
except StopAsyncIteration:
return
async def test_next_async_generator(self):
async def async_gen():
for item in slist:
yield item
it = ait.iter(async_gen())
self.assertEqual(await ait.next(it), "A")
self.assertEqual(await ait.next(it), "B")
self.assertEqual(await ait.next(it), "C")
with self.assertRaises(StopAsyncIteration):
await ait.next(it)
async def test_iter_iterator(self):
sentinel = object()
class async_iterator:
def __aiter__(self):
return sentinel
def __anext__(self):
return sentinel
aiter = async_iterator()
self.assertEqual(ait.iter(aiter), aiter)
async def wait(aiter):
aiter = aioitertools.iter(aiter)
items = []
next_future = asyncio.create_task(aioitertools.next(aiter))
while True:
if next_future.done() or not items:
try:
items.append(asyncio.create_task(await next_future))
next_future = asyncio.create_task(aioitertools.next(aiter))
except StopAsyncIteration:
break
await asyncio.wait([items[0], next_future],
return_when=asyncio.FIRST_COMPLETED)
while items and items[0].done():
yield items.pop(0).result()
for item in items:
yield await item
async def location_stream(
self,
ips: _IPsType,
*,
fields: _FieldsType = None,
lang: Optional[str] = None,
timeout: _TimeoutType = aiohttp.helpers.sentinel
) -> AsyncIterable[Dict[str, Any]]:
"""Returns async generator for locating IPs from iterable or async iterable
The method always uses batch API: https://ip-api.com/docs/api:batch
Parameters:
:param ips: The iterable or async iterable of IPs or dicts with additional info (see API docs)
:param fields: The sequence or set of returned fields in the result
:param lang: The language of the result
:param timeout: The timeout of the whole request to API
:return: async generator of results for every IP
"""
if self.closed:
raise ValueError('The client session is already closed')
if not isinstance(ips, (abc.Iterable, abc.AsyncIterable)):
raise TypeError(
"'ips' argument must be an iterable or async iterable")
if fields and not isinstance(fields, (abc.Sequence, abc.Set)):
raise TypeError("'fields' argument must be a sequence or set")
if lang and not isinstance(lang, str):
raise TypeError("'lang' argument must be a string")
fields = fields or self._fields
lang = lang or self._lang
url = self._make_url(config.batch_endpoint, fields, lang)
async for ips_batch in chunker(ips, chunk_size=config.batch_size):
results = await self._fetch_result(self._fetch_batch, url,
ips_batch, timeout)
async for result in aioitertools.iter(results):
yield result
async def collate(*iterables, key=lambda x: x):
iters = [aioitertools.iter(iterable) for iterable in iterables]
heap = []
nexts = []
for i, iterable in enumerate(iters):
try:
item = await aioitertools.next(iterable)
heapq.heappush(heap, (key(item), item, i))
except StopAsyncIteration:
nexts.append(None)
else:
nexts.append(asyncio.create_task(aioitertools.next(iterable)))
while any(fut is not None for fut in nexts):
k, item, i = heapq.heappop(heap)
yield item
if nexts[i] is not None:
try:
item = await nexts[i]
heapq.heappush(heap, (key(item), item, i))
except StopAsyncIteration:
nexts[i] = None
else:
nexts[i] = asyncio.create_task(aioitertools.next(iters[i]))
def __init__(self, iterator):
super().__init__()
self.iterator = iter(iterator)
self.cache = []
self.pointer = -1
async def _func(it):
it = aioitertools.iter(it)
async for item in it:
if query.predicate(item):
yield item
async def _func(it):
it = aioitertools.iter(it)
count = 0
async for _ in it:
count += 1
yield {"count": count}
async def _func(it):
it = aioitertools.iter(it)
yield {query.key: it}
async def _func(it):
it = aioitertools.iter(it)
async for item in it:
if item:
yield {k: item[k] for k in query.projector}
async def _func(it):
it = aioitertools.iter(it)
# islice takes start, stop
async for i in aioitertools.islice(it, query.skip, None):
yield i
async def _func(it):
it = aioitertools.iter(it)
# islice takes start, stop
async for i in aioitertools.islice(it, 0, query._count):
yield i
async def flatten(iterable):
"Flatten one level of nesting"
async for aiter in aioitertools.iter(iterable):
async for i in aioitertools.iter(aiter):
yield i
async def solve(
self
) -> Tuple[DefaultDict["libkol.Slot", Optional["libkol.Item"]],
Optional["libkol.Familiar"], List["libkol.Familiar"], ]:
from libkol import Modifier, Slot, Item, Familiar
# Get variables for some specific items
crown = await Item["Crown of Thrones"]
bjorn = await Item["Buddy Bjorn"]
# Load smithsness bonuses for tracking
smithsness_bonuses = {
s.item.id: await s.get_value()
async for s in
(Bonus.filter(modifier=Modifier.Smithsness,
item_id__not_isnull=True).prefetch_related("item"))
}
# Load hobo power bonuses for tracking
hobo_power_bonuses = {
s.item.id: await s.get_value()
async for s in
(Bonus.filter(modifier=Modifier.HoboPower,
item_id__not_isnull=True).prefetch_related("item"))
}
# Load relevant bonuses
modifiers = set(self.maximize + self.minimize +
list(self.maximum.keys()) + list(self.minimum.keys()))
bonuses = [
b async for b in (Bonus.filter(
effect_id__isnull=True, modifier__in=modifiers).filter(
Q(item_id__isnull=True)
| Q(item__hat=True)
| Q(item__shirt=True)
| Q(item__weapon=True)
| Q(item__offhand=True)
| Q(item__pants=True)
| Q(item__accessory=True)
| Q(item__familiar_equipment=True)).prefetch_related(
"familiar",
"item",
"outfit__variants__pieces",
"outfit__variants",
"outfit",
"throne_familiar",
))
]
grouped_bonuses = {} # type: Dict[libkol.Modifier, List[libkol.Bonus]]
for b in bonuses:
grouped_bonuses[b.modifier] = grouped_bonuses.get(b.modifier,
[]) + [b]
possible_items = (
[b.item for b in bonuses if isinstance(b.item, Item)] +
self.must_equip + [bjorn, crown])
possible_familiars = [
b.familiar for b in bonuses if isinstance(b.familiar, Familiar)
]
possible_throne_familiars = [
b.throne_familiar for b in bonuses
if isinstance(b.throne_familiar, Familiar)
]
# Define the problem
prob = LpProblem(self.summarise(), LpMaximize)
solution = LpVariable.dicts(
"outfit",
{repr(i)
for i in possible_items + possible_familiars}
| {self.enthroned_repr(f)
for f in possible_throne_familiars},
0,
3,
cat="Integer",
)
# Value of our Smithsness bonus
smithsness = self.calculate_smithsness(solution, smithsness_bonuses)
# Value of our Hobo Power bonus
hobo_power = self.calculate_hobo_power(solution, hobo_power_bonuses)
# Value of our familiar weight
familiar_weight = next(
(f.weight for f in possible_familiars if solution[repr(f)] == 1),
0)
# Objective
prob += lpSum([
m.sum([
await b.get_value(
smithsness=smithsness,
familiar_weight=familiar_weight,
hobo_power=hobo_power,
) *
(solution[repr(b.item)] if isinstance(b.item, Item) else 1) *
(solution[repr(b.familiar)] if isinstance(
b.familiar, Familiar) else 1) *
(solution[self.enthroned_repr(b.throne_familiar)]
if isinstance(b.throne_familiar, Familiar) else 1)
for b in bonuses if b.outfit is None or
(b.outfit and await b.outfit.is_fulfilled([
sb.item for sb in bonuses if isinstance(sb.item, Item)
and solution[repr(sb.item)] >= 1
]))
]) * self.weight[m] *
(1 if m in self.maximize else -1 if m in self.minimize else 0)
async for m, bonuses in iter(grouped_bonuses.items())
])
# Add minima and maxima
for m, bonuses in grouped_bonuses.items():
total = lpSum([
await b.get_value(smithsness=smithsness,
familiar_weight=familiar_weight) *
solution[repr(b.item)] for b in bonuses if b.item
])
if self.minimum.get(m) is not None:
prob += total >= self.minimum[m]
if self.maximum.get(m) is not None:
prob += total <= self.maximum[m]
# Maximum slot sizes
slot_sizes = [
("hat", 1),
("shirt", 1),
("weapon", 1),
("offhand", 1),
("pants", 1),
("accessory", 3),
("familiar_equipment", 1),
]
for slot, size in slot_sizes:
prob += (lpSum([
solution[repr(i)] for i in possible_items if getattr(i, slot)
]) <= size)
# Only use one familiar
prob += lpSum([solution[repr(f)] for f in possible_familiars]) <= 1
# Do not use familiars we don't have
for f in possible_familiars:
if f.have is False:
prob += solution[repr(f)] == 0
# Only throne familiars with throneable equips
prob += (lpSum([
solution[self.enthroned_repr(f)] for f in possible_throne_familiars
]) <= solution[repr(crown)] + solution[repr(bjorn)])
# Do not enthrone familiars we don't have
for f in possible_throne_familiars:
if f.have is False:
prob += solution[self.enthroned_repr(f)] == 0
# You've only got so many hands!
prob += (lpSum([
solution[repr(i)] for i in possible_items
if (i.weapon and i.weapon_hands >= 2) or i.offhand
]) <= 1)
# For each item...
for i in possible_items:
# Don't plan to equip things we can't wear
if i.meet_requirements() is False:
prob += solution[repr(i)] == 0
# We can only equip as many as we have
prob += solution[repr(i)] <= i.amount
# We can only equip one single equip item
if i.single_equip:
prob += solution[repr(i)] <= 1
# Forced equips
for i in self.must_equip:
prob += solution[repr(i)] >= 1
# Forced non-equips
for i in self.must_not_equip:
prob += solution[repr(i)] == 0
prob.writeLP("maximizer.lp")
prob.solve()
if prob.status is not LpStatusOptimal:
raise ValueError(LpStatus[prob.status])
familiar = None
throne_familiars = [] # type: List[Familiar]
result = defaultdict(
lambda: None) # type: DefaultDict[Slot, Optional[Item]]
for v in prob.variables():
index = v.name
q = v.varValue
index_parts = index.split("_")
if q == 0 or q is None or len(index_parts) < 3:
continue
id = int(index_parts[2])
if index_parts[1] == "<Familiar:":
familiar = next(f for f in possible_familiars if f.id == id)
continue
if index_parts[1] == "<Familiar(Enthroned):":
throne_familiars += next(f for f in possible_throne_familiars
if f.id == id)
continue
item = next(i for i in possible_items if i.id == id)
item_slot = item.slot # type: Slot
if item_slot == Slot.Acc1:
if result[Slot.Acc1] is None:
item_slot = Slot.Acc1
elif result[Slot.Acc2] is None:
item_slot = Slot.Acc2
elif result[Slot.Acc3] is None:
item_slot = Slot.Acc3
else:
raise Exception("Pulp has done something wrong")
result[item_slot] = item
return result, familiar, throne_familiars
async def matcher(d: dict) -> bool:
async for value in iter(d.values()):
if search_text in str(value).lower():
return True
return False
async def test_set(self):
self.assertEqual(await ait.set(ait.iter(slist)), set(slist))
