From c4310124e108c57a2beeefe9d1b1516b09ef9de6 Mon Sep 17 00:00:00 2001 From: Arijit Basu Date: Sat, 30 Dec 2023 10:38:47 +0530 Subject: [PATCH] Fix return type hints Examples: ```python result = await TaskPiccolo.select().where(TaskPiccolo.id == id).run() result[0][TaskPiccolo.name] result = ( await TaskPiccolo.insert(TaskPiccolo(name=name, completed=False)) .returning(TaskPiccolo.id) .run() ) result[0][TaskPiccolo.id] ``` --- piccolo/query/methods/insert.py | 3 ++- piccolo/query/methods/select.py | 6 +++--- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/piccolo/query/methods/insert.py b/piccolo/query/methods/insert.py index d7c655b80..a6ca68ced 100644 --- a/piccolo/query/methods/insert.py +++ b/piccolo/query/methods/insert.py @@ -18,7 +18,8 @@ class Insert( - t.Generic[TableInstance], Query[TableInstance, t.List[t.Dict[str, t.Any]]] + t.Generic[TableInstance], + Query[TableInstance, t.List[t.Dict[t.Union[str, "Column"], t.Any]]], ): __slots__ = ("add_delegate", "on_conflict_delegate", "returning_delegate") diff --git a/piccolo/query/methods/select.py b/piccolo/query/methods/select.py index a00745e48..4caafb8f9 100644 --- a/piccolo/query/methods/select.py +++ b/piccolo/query/methods/select.py @@ -279,7 +279,7 @@ def get_select_string( return f'SUM({column_name}) AS "{self._alias}"' -OptionalDict = t.Optional[t.Dict[str, t.Any]] +OptionalDict = t.Optional[t.Dict[t.Union[str, Column], t.Any]] class First(Proxy["Select", OptionalDict]): @@ -348,7 +348,7 @@ async def run( return dump_json(rows) -class Select(Query[TableInstance, t.List[t.Dict[str, t.Any]]]): +class Select(Query[TableInstance, t.List[t.Dict[t.Union[str, Column], t.Any]]]): __slots__ = ( "columns_list", "exclude_secrets", @@ -832,7 +832,7 @@ async def run( in_pool: bool = True, use_callbacks: bool = True, **kwargs, - ) -> t.List[t.Dict[str, t.Any]]: + ) -> t.List[t.Dict[t.Union[str, Column], t.Any]]: results = await super().run(node=node, in_pool=in_pool) if use_callbacks: return await self.callback_delegate.invoke(