def method_returning_dict(db: Any, *args: Any, **kwargs: Any) -> Any:
assert (returns is not None) # mypy bug
with get_cursor(db) as cur:
prepared_args = prepare_args_as_dict(func_def, args, kwargs)
cur.execute(hook(query.text, prepared_args), prepared_args)
names = [desc[0] for desc in cur.description]
if isinstance(returns.outer_dict_by, int):
keyidx = returns.outer_dict_by
else:
try:
keyidx = names.index(returns.outer_dict_by)
except ValueError:
raise KeyError(
f'key column {returns.outer_dict_by} not found')
if keyidx >= len(names):
raise IndexError(
f'key column index {keyidx} is out of range')
if returns.remove_key_column:
trimmed_names = names[0:keyidx] + names[keyidx + 1:]
process_row = generate_row_processor(
returns.inner_format, trimmed_names)
return {
row[keyidx]:
process_row(row[0:keyidx] + row[keyidx + 1:])
for row in cur
}
else:
process_row = generate_row_processor(
returns.inner_format, names)
return {row[keyidx]: process_row(row) for row in cur}
async def method_returning_single(db: Any, *args: Any, **kwargs: Any) -> Any:
assert(returns is not None) # mypy bug
async with get_cursor(db) as cur:
prepared_args = prepare_args_as_dict(func_def, args, kwargs)
await cur.execute(hook(query.text, prepared_args), prepared_args)
names = [desc[0] for desc in cur.description]
process_row = generate_row_processor(returns.inner_format, names)
return process_row(await cur.fetchone())
async def method_returning_none(db: Any, *args: Any,
**kwargs: Any) -> None:
prepared_args = prepare_args_as_dict(func_def, args, kwargs)
prepared_args_list = prepare_args_as_list(func_def, args, kwargs)
async with _get_connection(db) as conn:
await conn.execute(hook(query.text, prepared_args),
*prepared_args_list)
def test_dict_as_dict(func_def):
assert prepare_args_as_dict(func_def, ['a', 2, 'cc', 44], {}) == {
'a': 'a',
'b': 2,
'c': 'cc',
'd': 44,
'e': 'e',
'f': 6
}
def method_returning_list(db: Any, *args: Any,
**kwargs: Any) -> List[Any]:
assert (returns is not None) # mypy bug
with get_cursor(db) as cur:
prepared_args = prepare_args_as_dict(func_def, args, kwargs)
cur.execute(hook(query.text, prepared_args), prepared_args)
names = [desc[0] for desc in cur.description]
process_row = generate_row_processor(returns.inner_format,
names)
return [process_row(row) for row in cur]
async def method_returning_single(db: Any, *args: Any,
**kwargs: Any) -> Any:
assert (returns is not None) # mypy bug
prepared_args = prepare_args_as_dict(func_def, args, kwargs)
prepared_args_list = prepare_args_as_list(func_def, args, kwargs)
process_row = _generate_row_processor(returns.inner_format)
async with _get_connection(db) as conn:
return process_row(await conn.fetchrow(
hook(query.text, prepared_args), *prepared_args_list))
async def method_returning_iterator(
db: Any, *args: Any, **kwargs: Any) -> AsyncIterator[Any]:
assert (returns is not None) # mypy bug
prepared_args = prepare_args_as_dict(func_def, args, kwargs)
prepared_args_list = prepare_args_as_list(func_def, args, kwargs)
process_row = _generate_row_processor(returns.inner_format)
async with _get_connection(db, True) as conn:
async for row in conn.cursor(hook(query.text, prepared_args),
*prepared_args_list):
yield process_row(row)
async def method_returning_dict(db: Any, *args: Any,
**kwargs: Any) -> Any:
assert (returns is not None) # mypy bug
assert (returns.outer_dict_by is not None)
prepared_args = prepare_args_as_dict(func_def, args, kwargs)
prepared_args_list = prepare_args_as_list(func_def, args, kwargs)
if not returns.remove_key_column:
process_row = _generate_row_processor(returns.inner_format)
elif isinstance(returns.outer_dict_by, int):
process_row = _generate_row_processor_removing_index(
returns.inner_format, returns.outer_dict_by)
else:
process_row = _generate_row_processor_removing_name(
returns.inner_format, returns.outer_dict_by)
async with _get_connection(db) as conn:
return {
row[returns.outer_dict_by]: process_row(row)
for row in await conn.fetch(
hook(query.text, prepared_args), *prepared_args_list)
}
def method_returning_none(db: Any, *args: Any, **kwargs: Any) -> None:
with get_cursor(db) as cur:
prepared_args = prepare_args_as_dict(func_def, args, kwargs)
cur.execute(hook(query.text, prepared_args), prepared_args)
