From d523bf771e4899974326547554919e48475c51d9 Mon Sep 17 00:00:00 2001
From: Jakub Cierocki <jakub.cierocki@deepsense.ai>
Date: Wed, 3 Jul 2024 10:55:59 +0200
Subject: [PATCH] type hint fixes

---
 src/dbally/context/context.py         | 9 +++++----
 src/dbally/iql/_processor.py          | 1 -
 src/dbally/views/exposed_functions.py | 4 ++--
 3 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/src/dbally/context/context.py b/src/dbally/context/context.py
index 3e991aae..d3af935d 100644
--- a/src/dbally/context/context.py
+++ b/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
diff --git a/src/dbally/iql/_processor.py b/src/dbally/iql/_processor.py
index 49ab1eb9..b6a1648a 100644
--- a/src/dbally/iql/_processor.py
+++ b/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)
diff --git a/src/dbally/views/exposed_functions.py b/src/dbally/views/exposed_functions.py
index f133a973..c6d400d2 100644
--- a/src/dbally/views/exposed_functions.py
+++ b/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: