type hint fixes

Author: ds-jakub-cierocki

src/dbally/context/context.py

@@ -1,13 +1,14 @@
 import ast
 
-from typing import Optional, Type, Sequence
+from typing import Sequence, TypeVar
 from typing_extensions import Self
 from pydantic import BaseModel
 
 from dbally.context.exceptions import ContextNotAvailableError
 
 
-CustomContextsList = Sequence[Type['BaseCallerContext']]  # TODO confirm the naming
+CustomContext = TypeVar('CustomContext', bound='BaseCallerContext', covariant=True)
+CustomContextsList = Sequence[CustomContext]  # TODO confirm the naming
 
 
 class BaseCallerContext(BaseModel):
@@ -17,11 +18,11 @@ class BaseCallerContext(BaseModel):
     """
 
     @classmethod
-    def select_context(cls, contexts: Sequence[Type[Self]]) -> Type[Self]:
+    def select_context(cls, contexts: CustomContextsList) -> Self:
         if not contexts:
             raise ContextNotAvailableError("The LLM detected that the context is required to execute the query and the filter signature allows contextualization while the context was not provided.")
 
-        # we assume here that each element of `contexts` represents a different subclass of BaseCallerContext
+        # this method is called from the subclass of BaseCallerContext pointing the right type of custom context
         return next(filter(lambda obj: isinstance(obj, cls), contexts))
 
     @classmethod

src/dbally/iql/_processor.py

@@ -132,7 +132,6 @@ def _parse_arg(
                 raise ContextualisationNotAllowed("The LLM detected that the context is required to execute the query while the filter signature does not allow it at all.")
 
             if not _does_arg_allow_context(arg_spec):
-                print(arg_spec)
                 raise ContextualisationNotAllowed(f"The LLM detected that the context is required to execute the query while the filter signature does allow it for `{arg_spec.name}` argument.")
 
             return parent_func_def.context_class.select_context(self.contexts)

src/dbally/views/exposed_functions.py

@@ -1,7 +1,7 @@
 import re
 from dataclasses import dataclass
 from typing import _GenericAlias  # type: ignore
-from typing import List, Optional, Union, Type
+from typing import Sequence, Optional, Union, Type
 
 from dbally.similarity import AbstractSimilarityIndex
 from dbally.context.context import BaseCallerContext
@@ -57,7 +57,7 @@ class ExposedFunction:
 
     name: str
     description: str
-    parameters: List[MethodParamWithTyping]
+    parameters: Sequence[MethodParamWithTyping]
     context_class: Optional[Type[BaseCallerContext]] = None
 
     def __str__(self) -> str: