simplify structure of parsed AST

Author: Computerdores

tagstudio/src/core/query_lang/ast.py

@@ -1,6 +1,6 @@
 from abc import ABC, abstractmethod
 from enum import Enum
-from typing import Generic, TypeVar, Union
+from typing import Generic, TypeVar
 
 
 class ConstraintType(Enum):  # TODO add remaining ones
@@ -37,19 +37,19 @@ def __repr__(self) -> str:
 
 
 class ANDList(AST):
-    terms: list[Union["ORList", "Constraint", "Not"]]
+    terms: list[AST]
 
-    def __init__(self, terms: list[Union["ORList", "Constraint", "Not"]]) -> None:
+    def __init__(self, terms: list[AST]) -> None:
         super().__init__()
         for term in terms:
             term.parent = self
         self.terms = terms
 
 
 class ORList(AST):
-    elements: list[ANDList]
+    elements: list[AST]
 
-    def __init__(self, elements: list[ANDList]) -> None:
+    def __init__(self, elements: list[AST]) -> None:
         super().__init__()
         for element in elements:
             element.parent = self

tagstudio/src/core/query_lang/parser.py

@@ -1,5 +1,3 @@
-from typing import Union
-
 from src.core.query_lang.ast import AST, ANDList, Constraint, Not, ORList, Property
 from src.core.query_lang.tokenizer import ConstraintType, Token, Tokenizer, TokenType
 from src.core.query_lang.util import ParsingError
@@ -25,24 +23,24 @@ def parse(self) -> AST:
             raise ParsingError(self.next_token.start, self.next_token.end, "Syntax Error")
         return out
 
-    def __or_list(self) -> ORList:
+    def __or_list(self) -> AST:
         terms = [self.__and_list()]
 
         while self.__is_next_or():
             self.__eat(TokenType.ULITERAL)
             terms.append(self.__and_list())
 
-        return ORList(terms)
+        return ORList(terms) if len(terms) > 1 else terms[0]
 
     def __is_next_or(self) -> bool:
         return self.next_token.type == TokenType.ULITERAL and self.next_token.value.upper() == "OR"
 
-    def __and_list(self) -> ANDList:
+    def __and_list(self) -> AST:
         elements = [self.__term()]
         while self.next_token.type != TokenType.EOF and not self.__is_next_or():
             self.__skip_and()
             elements.append(self.__term())
-        return ANDList(elements)
+        return ANDList(elements) if len(elements) > 1 else elements[0]
 
     def __skip_and(self) -> None:
         if self.__is_next_and():
@@ -54,7 +52,7 @@ def __skip_and(self) -> None:
     def __is_next_and(self) -> bool:
         return self.next_token.type == TokenType.ULITERAL and self.next_token.value.upper() == "AND"
 
-    def __term(self) -> Union[ORList, Constraint, Not]:
+    def __term(self) -> AST:
         if self.__is_next_not():
             self.__eat(TokenType.ULITERAL)
             return Not(self.__term())
@@ -108,9 +106,3 @@ def __eat(self, type: TokenType) -> Token:
 
     def __syntax_error(self, msg: str = "Syntax Error") -> ParsingError:
         return ParsingError(self.next_token.start, self.next_token.end, msg)
-
-
-if __name__ == "__main__":  # TODO remove
-    print("")  # noqa: T201
-    p = Parser("Mario AND Luigi tag:test[parent=Color,color=red] OR mediatype:test")
-    print(p.parse())  # noqa: T201