From c76368cac525f343d0dbcda8f2357be6e4bb349d Mon Sep 17 00:00:00 2001
From: breathe <breathe@users.noreply.github.com>
Date: Thu, 16 May 2024 06:41:26 -0600
Subject: [PATCH] Improve type inference when using decorators with pyright
 (#172)

* Improve type inference when using decorators with pyright

* linting

---------

Co-authored-by: N. Ben Cohen <ben.cohen@zephr.xyz>
Co-authored-by: Brian Pugh <bnp117@gmail.com>
---
 cyclopts/core.py | 18 +++++++++++-------
 1 file changed, 11 insertions(+), 7 deletions(-)

diff --git a/cyclopts/core.py b/cyclopts/core.py
index 6a2e8246..9ffd0012 100644
--- a/cyclopts/core.py
+++ b/cyclopts/core.py
@@ -16,6 +16,7 @@
     Literal,
     Optional,
     Tuple,
+    TypeVar,
     Union,
 )
 
@@ -51,6 +52,9 @@
     to_tuple_converter,
 )
 
+T = TypeVar("T", bound=Callable)
+
+
 if sys.version_info < (3, 9):
     from typing_extensions import Annotated
 else:
@@ -446,10 +450,10 @@ def parse_commands(self, tokens: Union[None, str, Iterable[str]] = None):
 
     def command(
         self,
-        obj: Optional[Callable] = None,
+        obj: Optional[T] = None,
         name: Union[None, str, Iterable[str]] = None,
         **kwargs,
-    ) -> Callable:
+    ) -> T:
         """Decorator to register a function as a CLI command.
 
         Parameters
@@ -466,7 +470,7 @@ def command(
             Any argument that :class:`App` can take.
         """
         if obj is None:  # Called ``@app.command(...)``
-            return partial(self.command, name=name, **kwargs)
+            return partial(self.command, name=name, **kwargs)  # pyright: ignore[reportReturnType]
 
         if isinstance(obj, App):
             app = obj
@@ -506,18 +510,18 @@ def command(
 
         app._parents.append(self)
 
-        return obj
+        return obj  # pyright: ignore[reportReturnType]
 
     def default(
         self,
-        obj: Optional[Callable] = None,
+        obj: Optional[T] = None,
         *,
         converter=None,
         validator=None,
-    ):
+    ) -> T:
         """Decorator to register a function as the default action handler."""
         if obj is None:  # Called ``@app.default_command(...)``
-            return partial(self.default, converter=converter, validator=validator)
+            return partial(self.default, converter=converter, validator=validator)  # pyright: ignore[reportReturnType]
 
         if isinstance(obj, App):  # Registering a sub-App
             raise TypeError("Cannot register a sub-App to default.")