fix type checking for lat-bound callables

Author: apalala

Diff for: late/__init__.py

@@ -7,6 +7,7 @@
 __all__ = ['latebinding', 'late', '__']
 
 
+_R = TypeVar('_R')
 _T = TypeVar('_T')
 _V = TypeVar('_V')
 
@@ -15,7 +16,7 @@ class _LateBound(NamedTuple):
     actual: Any
 
 
-def late(o: _T | Iterator[_V]) -> _T | _V:
+def late(o: _T | Iterator[_V] | Callable[[], _R]) -> _T | _V | _R:
     if isinstance(o, int | float | str | bool | bytes | bytearray | frozenset):
         return o  # type: ignore
 

Diff for: test/complex_test.py

@@ -23,7 +23,7 @@ def a() -> int:
         return t
 
     @latebinding
-    def f(x: int = __(a)) -> int:  # type: ignore
+    def f(x: int = __(a)) -> int:
         return 2 * x
 
     assert f() == 2