type hint lifting

Author: ds-jakub-cierocki

src/dbally/collection/collection.py

@@ -16,7 +16,7 @@
 from dbally.similarity.index import AbstractSimilarityIndex
 from dbally.view_selection.base import ViewSelector
 from dbally.views.base import BaseView, IndexLocation
-from dbally.context.context import BaseCallerContext
+from dbally.context.context import BaseCallerContext, CustomContextsList
 
 
 class Collection:
@@ -157,7 +157,7 @@ async def ask(
         dry_run: bool = False,
         return_natural_response: bool = False,
         llm_options: Optional[LLMOptions] = None,
-        context: Optional[List[BaseCallerContext]] = None
+        context: Optional[CustomContextsList] = None
     ) -> ExecutionResult:
         """
         Ask question in a text form and retrieve the answer based on the available views.

src/dbally/context/_utils.py

@@ -1,18 +1,15 @@
-import typing
 import typing_extensions as type_ext
 
+from typing import Sequence, Tuple, Optional, Type, Any, Union
 from inspect import isclass
 
 from dbally.context.context import BaseCallerContext
 from dbally.views.exposed_functions import MethodParamWithTyping
 
 
 def _extract_params_and_context(
-    filter_method_: typing.Callable, hidden_args: typing.List[str]
-) -> typing.Tuple[
-    typing.List[MethodParamWithTyping],
-    typing.Optional[typing.Type[BaseCallerContext]]
-]:
+    filter_method_: type_ext.Callable, hidden_args: Sequence[str]
+) -> Tuple[Sequence[MethodParamWithTyping], Optional[Type[BaseCallerContext]]]:
     """
     Processes the MethodsBaseView filter method signauture to extract the args and type hints in the desired format.
     Context claases are getting excluded the returned MethodParamWithTyping list. Only the first BaseCallerContext
@@ -29,7 +26,7 @@ class is returned.
     params = []
     context = None
     # TODO confirm whether to use typing.get_type_hints(method) or method.__annotations__
-    for name_, type_ in typing.get_type_hints(filter_method_).items():
+    for name_, type_ in type_ext.get_type_hints(filter_method_).items():
         if name_ in hidden_args:
             continue
 
@@ -39,11 +36,11 @@ class is returned.
             # this is the case when user provides a context but no other type hint for a specifc arg
             # TODO confirm whether this case should be supported
             context = type_
-            type_ = typing.Any
-        elif type_ext.get_origin(type_) is typing.Union:
+            type_ = Any
+        elif type_ext.get_origin(type_) is Union:
             union_subtypes = type_ext.get_args(type_)
             if not union_subtypes:
-                type_ = typing.Any
+                type_ = Any
 
             for subtype_ in union_subtypes:  # type: ignore
                 # TODO add custom error for the situation when user provides more than two contexts for a single filter
@@ -58,7 +55,7 @@ class is returned.
 
 
 def _does_arg_allow_context(arg: MethodParamWithTyping) -> bool:
-    if type_ext.get_origin(arg.type) is not typing.Union and not issubclass(arg.type, BaseCallerContext):
+    if type_ext.get_origin(arg.type) is not Union and not issubclass(arg.type, BaseCallerContext):
         return False
 
     for subtype in type_ext.get_args(arg.type):

src/dbally/context/context.py

@@ -1,14 +1,13 @@
 import ast
 
-from typing import List, Optional, Type, TypeVar
+from typing import Optional, Type, Sequence
 from typing_extensions import Self
 from pydantic import BaseModel
 
 from dbally.context.exceptions import ContextNotAvailableError
 
 
-T = TypeVar('T', bound='BaseCallerContext')
-AllCallerContexts = Optional[List[T]]  # TODO confirm the naming
+CustomContextsList = Sequence[Type['BaseCallerContext']]  # TODO confirm the naming
 
 
 class BaseCallerContext(BaseModel):
@@ -18,7 +17,7 @@ class BaseCallerContext(BaseModel):
     """
 
     @classmethod
-    def select_context(cls, contexts: List[T]) -> T:
+    def select_context(cls, contexts: Sequence[Type[Self]]) -> Type[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.")
 

src/dbally/iql/_processor.py

@@ -1,7 +1,6 @@
 import ast
 
-from typing import TYPE_CHECKING, Any, List, Optional, Union, Mapping, Dict
-from typing_extensions import Callable
+from typing import TYPE_CHECKING, Any, List, Optional, Union, Mapping, Type
 
 from dbally.audit.event_tracker import EventTracker
 from dbally.iql import syntax
@@ -13,7 +12,7 @@
     IQLUnsupportedSyntaxError,
 )
 from dbally.iql._type_validators import validate_arg_type
-from dbally.context.context import BaseCallerContext
+from dbally.context.context import BaseCallerContext, CustomContextsList
 from dbally.context.exceptions import ContextNotAvailableError, ContextualisationNotAllowed
 from dbally.context._utils import _extract_params_and_context, _does_arg_allow_context
 from dbally.views.exposed_functions import MethodParamWithTyping, ExposedFunction
@@ -25,15 +24,15 @@ class IQLProcessor:
     """
     source: str
     allowed_functions: Mapping[str, "ExposedFunction"]
-    contexts: List[BaseCallerContext]
+    contexts: CustomContextsList
     _event_tracker: EventTracker
 
 
     def __init__(
         self,
         source: str,
         allowed_functions: List["ExposedFunction"],
-        contexts: Optional[List[BaseCallerContext]] = None,
+        contexts: Optional[CustomContextsList] = None,
         event_tracker: Optional[EventTracker] = None
     ) -> None:
         self.source = source
@@ -52,7 +51,7 @@ async def process(self) -> syntax.Node:
         Raises:
              IQLError: if parsing fails.
         """
-        # TODO adjust this method to prevent making context class constructor calls lowercase
+
         self.source = self._to_lower_except_in_quotes(self.source, ["AND", "OR", "NOT"])
 
         ast_tree = ast.parse(self.source)
@@ -132,15 +131,11 @@ def _parse_arg(
             if parent_func_def.context_class is None:
                 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 _does_arg_allow_context(arg_spec):
+            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.")
 
-            context = parent_func_def.context_class.select_context(self.contexts)
-
-            try:
-                return getattr(context, arg_spec.name)
-            except AttributeError:
-                raise ContextNotAvailableError(f"The LLM detected that the context is required to execute the query and the context object was provided but it is missing the `{arg_spec.name}` field.")
+            return parent_func_def.context_class.select_context(self.contexts)
 
         if not isinstance(arg, ast.Constant):
             raise IQLArgumentParsingError(arg, self.source)

src/dbally/iql/_query.py

@@ -1,9 +1,9 @@
-from typing import TYPE_CHECKING, List, Optional
+from typing import TYPE_CHECKING, List, Optional, Type
 
 from ..audit.event_tracker import EventTracker
 from . import syntax
 from ._processor import IQLProcessor
-from dbally.context.context import BaseCallerContext
+from dbally.context.context import BaseCallerContext, CustomContextsList
 
 if TYPE_CHECKING:
     from dbally.views.structured import ExposedFunction
@@ -25,7 +25,7 @@ async def parse(
         source: str,
         allowed_functions: List["ExposedFunction"],
         event_tracker: Optional[EventTracker] = None,
-        context: Optional[List[BaseCallerContext]] = None
+        context: Optional[CustomContextsList] = None
     ) -> "IQLQuery":
         """
         Parse IQL string to IQLQuery object.

src/dbally/views/base.py

@@ -1,12 +1,12 @@
 import abc
-from typing import Dict, List, Optional, Tuple
+from typing import Dict, List, Optional, Tuple, Type
 
 from dbally.audit.event_tracker import EventTracker
 from dbally.collection.results import ViewExecutionResult
 from dbally.llms.base import LLM
 from dbally.llms.clients.base import LLMOptions
 from dbally.similarity import AbstractSimilarityIndex
-from dbally.context.context import BaseCallerContext
+from dbally.context.context import BaseCallerContext, CustomContextsList
 
 IndexLocation = Tuple[str, str, str]
 
@@ -26,7 +26,7 @@ async def ask(
         n_retries: int = 3,
         dry_run: bool = False,
         llm_options: Optional[LLMOptions] = None,
-        context: Optional[List[BaseCallerContext]] = None
+        context: Optional[CustomContextsList] = None
     ) -> ViewExecutionResult:
         """
         Executes the query and returns the result.

src/dbally/views/structured.py

@@ -1,6 +1,6 @@
 import abc
 from collections import defaultdict
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, Type
 
 from dbally.audit.event_tracker import EventTracker
 from dbally.collection.results import ViewExecutionResult
@@ -10,7 +10,7 @@
 from dbally.llms.base import LLM
 from dbally.llms.clients.base import LLMOptions
 from dbally.views.exposed_functions import ExposedFunction
-from dbally.context.context import BaseCallerContext
+from dbally.context.context import BaseCallerContext, CustomContextsList
 
 from ..similarity import AbstractSimilarityIndex
 from .base import BaseView, IndexLocation
@@ -42,7 +42,7 @@ async def ask(
         n_retries: int = 3,
         dry_run: bool = False,
         llm_options: Optional[LLMOptions] = None,
-        context: Optional[List[BaseCallerContext]] = None
+        context: Optional[CustomContextsList] = None
     ) -> ViewExecutionResult:
         """
         Executes the query and returns the result. It generates the IQL query from the natural language query\