Improve filter_tokens in core tokenizer

nazargesyk · nazargesyk · commit b40132aa1ba8 · 2024-07-17T11:29:12.000+03:00
diff --git a/uncoder-core/app/translator/core/tokenizer.py b/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))
