Improve typing

.github/pull_request_template.md

@@ -0,0 +1,11 @@
+## Description
+
+<!-- Describe the goal of the PR in a few lines. -->
+
+## TODO
+
+- [ ] Update version.
+- [ ] Update [changelog](../CHANGELOG.md).
+- [ ] Quality passed.
+- [ ] Tests passed.
+- [ ] Readme is up-to-date.

CHANGELOG.md

@@ -7,8 +7,15 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Changed
+
+- Typing has been improved.
+
+## [0.1.0]
+
 ### Added
 
 - Various observations contexts to track code.
 
-[Unreleased]: https://github.com/bendabir/trackotron/tree/develop
+[Unreleased]: https://github.com/bendabir/trackotron/compare/0.1.0...main
+[0.1.0]: https://github.com/bendabir/trackotron/tree/0.1.0

README.md

@@ -117,3 +117,4 @@ Here are some ideas to improve the current decorator :
 - [ ] Add more tests to ensure behaviors and consistency
 - [ ] Support context reuse if that makes sense
 - [ ] Handle data obfuscation when needed so we don't expose secrets.
+- [ ] Make parenthesis usage optional (see [#918](https://github.com/langfuse/langfuse-python/pull/918))

pyproject.toml

@@ -1,7 +1,7 @@
 # --- poetry ---
 [tool.poetry]
 name = "trackotron"
-version = "0.1.0"
+version = "0.2.0"
 description = "Improved decorator for Langfuse."
 license = "GPL-3.0-or-later"
 authors = ["Bendabir"]
@@ -37,13 +37,13 @@ build-backend = "poetry.core.masonry.api"
 # --- black ---
 [tool.black]
 line-length = 88
-target-version = ["py38"]
+target-version = ["py39"]
 preview = false
 
 # --- mypy ---
 [tool.mypy]
 local_partial_types = true
-python_version = "3.8"
+python_version = "3.9"
 mypy_path = "src/,tests/"
 files = "**/src/**/*.py,**/tests/**/*.py"
 scripts_are_modules = true
@@ -79,15 +79,10 @@ enable_error_code = [
     "unused-awaitable",
 ]
 
-[tool.pydantic-mypy]
-init_typed = false
-init_forbid_extra = true
-warn_required_dynamic_aliases = true
-
 # --- ruff ---
 [tool.ruff]
 line-length = 88
-target-version = "py38"
+target-version = "py39"
 indent-width = 4
 preview = true
 namespace-packages = ["src/trackotron"]

src/trackotron/contexts/context.py

@@ -12,7 +12,6 @@
     TYPE_CHECKING,
     Any,
     Callable,
-    Coroutine,
     Generic,
     TypeVar,
     final,
@@ -28,9 +27,18 @@
     from contextvars import ContextVar
     from types import TracebackType
 
-    from trackotron.types_ import Arguments, P, R_co
+    from trackotron.types_ import (
+        Arguments,
+        AsyncFunction,
+        AsyncInjectedFunction,
+        Function,
+        InjectedFunction,
+        ObservationParameters,
+        P,
+        R_co,
+        TraceParameters,
+    )
     from trackotron.types_.compatibility import Langfuse, StatefulClient
-    from trackotron.types_.misc import ObservationParameters, TraceParameters
 
 U_co = TypeVar("U_co", bound=ObservationUpdate, covariant=True)
 
@@ -238,12 +246,9 @@ def _is_method(f: Callable[P, R_co], /) -> bool:
     @final
     def _wrap_sync(
         self,
-        f: Callable[
-            Concatenate[ObservationProxy[O_co, U_co], P],
-            R_co,
-        ],
+        f: InjectedFunction[ObservationProxy[O_co, U_co], P, R_co],
         /,
-    ) -> Callable[P, R_co]:
+    ) -> Function[P, R_co]:
         @ft.wraps(f)
         def wrap(*args: Any, **kwargs: Any) -> R_co:
             with self._context() as proxy:
@@ -288,43 +293,31 @@ def wrap(*args: Any, **kwargs: Any) -> R_co:
     @final
     def _wrap_async(
         self,
-        f: Callable[
-            Concatenate[ObservationProxy[O_co, U_co], P],
-            Coroutine[Any, Any, R_co],
-        ],
+        f: AsyncInjectedFunction[ObservationProxy[O_co, U_co], P, R_co],
         /,
-    ) -> Callable[P, Coroutine[Any, Any, R_co]]:
+    ) -> AsyncFunction[P, R_co]:
         raise NotImplementedError
 
     @overload
     def __call__(
         self,
-        f: Callable[
-            Concatenate[ObservationProxy[O_co, U_co], P],
-            R_co,
-        ],
+        f: InjectedFunction[ObservationProxy[O_co, U_co], P, R_co],
         /,
-    ) -> Callable[P, R_co]: ...
+    ) -> Function[P, R_co]: ...
 
     @overload
     def __call__(
         self,
-        f: Callable[
-            Concatenate[ObservationProxy[O_co, U_co], P],
-            Coroutine[Any, Any, R_co],
-        ],
+        f: AsyncInjectedFunction[ObservationProxy[O_co, U_co], P, R_co],
         /,
-    ) -> Callable[P, Coroutine[Any, Any, R_co]]: ...
+    ) -> AsyncFunction[P, R_co]: ...
 
     @final
     def __call__(
         self,
-        f: Callable[
-            Concatenate[ObservationProxy[O_co, U_co], P],
-            R_co,
-        ],
+        f: InjectedFunction[ObservationProxy[O_co, U_co], P, R_co],
         /,
-    ) -> Callable[P, R_co] | Callable[P, Coroutine[Any, Any, R_co]]:
+    ) -> Function[P, R_co] | AsyncFunction[P, R_co]:
         """Decorate a function to observe it in Langfuse.
 
         If the name was not provided, it will infer it from the function/method.

src/trackotron/types_/__init__.py

@@ -2,6 +2,20 @@
 
 from __future__ import annotations
 
+from .aliases import (
+    AsyncClassMethod,
+    AsyncFunction,
+    AsyncInjectedClassMethod,
+    AsyncInjectedFunction,
+    AsyncInjectedMethod,
+    AsyncMethod,
+    ClassMethod,
+    Function,
+    InjectedClassMethod,
+    InjectedFunction,
+    InjectedMethod,
+    Method,
+)
 from .compatibility import (
     Langfuse,
     ModelUsage,
@@ -13,12 +27,26 @@
     StatefulSpanClient,
     StatefulTraceClient,
 )
-from .generics import P, R_co, T_contra
+from .generics import C, I, P, R_co
 from .misc import Arguments, ObservationParameters, ObservationType, TraceParameters
 
 __all__ = [
     "Arguments",
+    "AsyncClassMethod",
+    "AsyncFunction",
+    "AsyncInjectedClassMethod",
+    "AsyncInjectedFunction",
+    "AsyncInjectedMethod",
+    "AsyncMethod",
+    "C",
+    "ClassMethod",
+    "Function",
+    "I",
+    "InjectedClassMethod",
+    "InjectedFunction",
+    "InjectedMethod",
     "Langfuse",
+    "Method",
     "ModelUsage",
     "ObservationParameters",
     "ObservationType",
@@ -31,6 +59,5 @@
     "StatefulGenerationClient",
     "StatefulSpanClient",
     "StatefulTraceClient",
-    "T_contra",
     "TraceParameters",
 ]

src/trackotron/types_/aliases.py

@@ -0,0 +1,26 @@
+"""Copyright (c) 2024 Bendabir."""
+
+from __future__ import annotations
+
+from typing import Any, Callable, Coroutine
+
+from typing_extensions import Concatenate
+
+from trackotron.types_.generics import C, I, P, R_co
+
+Function = Callable[P, R_co]
+InjectedFunction = Callable[Concatenate[I, P], R_co]
+Method = Callable[Concatenate[C, P], R_co]
+InjectedMethod = Callable[Concatenate[C, I, P], R_co]
+ClassMethod = Callable[Concatenate[type[C], P], R_co]
+InjectedClassMethod = Callable[Concatenate[type[C], I, P], R_co]
+
+AsyncFunction = Callable[P, Coroutine[Any, Any, R_co]]
+AsyncInjectedFunction = Callable[Concatenate[I, P], Coroutine[Any, Any, R_co]]
+AsyncMethod = Callable[Concatenate[C, P], Coroutine[Any, Any, R_co]]
+AsyncInjectedMethod = Callable[Concatenate[C, I, P], Coroutine[Any, Any, R_co]]
+AsyncClassMethod = Callable[Concatenate[type[C], P], Coroutine[Any, Any, R_co]]
+AsyncInjectedClassMethod = Callable[
+    Concatenate[type[C], I, P],
+    Coroutine[Any, Any, R_co],
+]

src/trackotron/types_/generics.py

@@ -8,4 +8,5 @@
 
 P = ParamSpec("P")
 R_co = TypeVar("R_co", covariant=True)
-T_contra = TypeVar("T_contra", contravariant=True)
+C = TypeVar("C")
+I = TypeVar("I")