From 846b3ab36b81b05685bd82aec34a4be687e38371 Mon Sep 17 00:00:00 2001 From: David Rajaratnam Date: Mon, 4 Mar 2024 11:46:25 +1100 Subject: [PATCH] Added overload signature for query select with a function --- clorm/orm/_queryimpl.py | 9 +++++++++ clorm/orm/query.py | 2 +- 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/clorm/orm/_queryimpl.py b/clorm/orm/_queryimpl.py index 147d73a..6700d75 100644 --- a/clorm/orm/_queryimpl.py +++ b/clorm/orm/_queryimpl.py @@ -736,6 +736,11 @@ def group_by(self, *expressions: Any) -> "GroupedQuery[Any, _T]": # -------------------------------------------------------------------------- # Explicitly select the elements to output or delete # -------------------------------------------------------------------------- + + @overload + def select(self, __ent0: Callable[..., _T0]) -> "UnGroupedQuery[_T0]": + ... + # START OVERLOADED FUNCTIONS self.select;UnGroupedQuery[{0}];1;5;Type;Y # code within this block is **programmatically, @@ -1153,6 +1158,10 @@ def group_by(self, *expressions): # doesn't need to implement group_by raise ValueError("Cannot specify 'group_by' multiple times") + @overload + def select(self, __ent0: Callable[..., _T0]) -> "GroupedQuery[_KT, _T0]": + ... + # START OVERLOADED FUNCTIONS self.select;GroupedQuery[_KT, {0}];1;5;Type;Y # code within this block is **programmatically, diff --git a/clorm/orm/query.py b/clorm/orm/query.py index a6b5c5e..624a41a 100644 --- a/clorm/orm/query.py +++ b/clorm/orm/query.py @@ -230,7 +230,7 @@ def functor(self): def func(paths: List[Any], functor: Callable[..., _T]) -> _T: """Return a wrapped functor for use in a 'query' where or 'select' clause""" - return FuncInputSpec(paths, functor) + return cast(_T, FuncInputSpec(paths, functor)) # Note: used to just return a FunctionComparator but that is only suitable # for a where clause