Merge pull request #174 from UncoderIO/gis-8171

Improve filter_tokens in core tokenizer

Author: nazargesyk

uncoder-core/app/translator/core/parser.py

@@ -28,6 +28,7 @@
 from app.translator.core.models.platform_details import PlatformDetails
 from app.translator.core.models.query_container import RawQueryContainer, TokenizedQueryContainer
 from app.translator.core.models.query_tokens.field import Field
+from app.translator.core.models.query_tokens.field_field import FieldField
 from app.translator.core.models.query_tokens.field_value import FieldValue
 from app.translator.core.models.query_tokens.function_value import FunctionValue
 from app.translator.core.tokenizer import QueryTokenizer
@@ -68,6 +69,11 @@ def get_field_tokens(
         for token in query_tokens:
             if isinstance(token, FieldValue):
                 field_tokens.append(token.field)
+            elif isinstance(token, FieldField):
+                if token.field_left:
+                    field_tokens.append(token.field_left)
+                if token.field_right:
+                    field_tokens.append(token.field_right)
             elif isinstance(token, FunctionValue):
                 field_tokens.extend(self.tokenizer.get_field_tokens_from_func_args([token.function]))
 

uncoder-core/app/translator/core/tokenizer.py

@@ -338,7 +338,7 @@ def tokenize(self, query: str) -> list[QUERY_TOKEN_TYPE]:
     @staticmethod
     def filter_tokens(
         tokens: list[QUERY_TOKEN_TYPE],
-        token_type: Union[type[FieldValue], type[Field], type[Keyword], type[Identifier]],
+        token_type: Union[type[FieldValue], type[Field], type[FieldField], type[Keyword], type[Identifier]],
     ) -> list[QUERY_TOKEN_TYPE]:
         return [token for token in tokens if isinstance(token, token_type)]
 
@@ -363,7 +363,9 @@ def get_field_tokens_from_func_args(  # noqa: PLR0912
                 result.extend(self.get_field_tokens_from_func_args(args=arg.args))
                 result.extend(self.get_field_tokens_from_func_args(args=arg.by_clauses))
                 result.extend(self.get_field_tokens_from_func_args(args=[arg.filter_]))
-            elif isinstance(arg, (JoinFunction, UnionFunction)):
+            elif isinstance(arg, JoinFunction):
+                result.extend(self.get_field_tokens_from_func_args(args=arg.condition))
+            elif isinstance(arg, UnionFunction):
                 continue
             elif isinstance(arg, Function):
                 result.extend(self.get_field_tokens_from_func_args(args=arg.args))