query: move selects logic to subclasses

Author: Gerit Wagner

search_query/ebsco/parser.py

@@ -13,6 +13,7 @@
 from search_query.parser_base import QueryStringParser
 from search_query.query import Query
 from search_query.query import SearchField
+from search_query.query_near import NEARQuery
 from search_query.query_term import Term
 
 
@@ -244,11 +245,12 @@ def parse_query_tree(
 
             elif token.type == TokenTypes.PROXIMITY_OPERATOR:
                 distance = self._extract_proximity_distance(token)
-                proximity_node = Query(
+                proximity_node = NEARQuery(
                     value=token.value,
+                    distance=distance,
+                    children=[],
                     position=token.position,
                     search_field=search_field or field_context,
-                    distance=distance,
                     platform="deactivated",
                 )
                 parent, current_operator = self.append_operator(parent, proximity_node)

search_query/ebsco/serializer.py

@@ -13,6 +13,8 @@
 
 def to_string_ebsco(query: Query) -> str:
     """Convert the query to a string representation for EBSCO."""
+    # pylint: disable=import-outside-toplevel
+    from search_query.query_near import NEARQuery
 
     query = query.copy()
 
@@ -27,7 +29,7 @@ def to_string_ebsco(query: Query) -> str:
     for i, child in enumerate(query.children):
         child_str = to_string_ebsco(child)
 
-        if query.value in {"NEAR", "WITHIN"}:
+        if isinstance(query, NEARQuery):
             # Convert proximity operator to EBSCO format
             proximity_operator = (
                 f"{'N' if query.value == 'NEAR' else 'W'}{query.distance}"

search_query/query.py

@@ -3,10 +3,8 @@
 from __future__ import annotations
 
 import copy
-import re
 import typing
 
-from search_query.constants import Fields
 from search_query.constants import Operators
 from search_query.constants import PLATFORM
 from search_query.constants import SearchField
@@ -31,18 +29,15 @@ def __init__(
         search_field: typing.Optional[SearchField] = None,
         children: typing.Optional[typing.List[typing.Union[str, Query]]] = None,
         position: typing.Optional[typing.Tuple[int, int]] = None,
-        distance: typing.Optional[int] = None,
         platform: str = "generic",
     ) -> None:
         self._value: str = ""
         self._operator = False
-        self._distance = None
         self._children: typing.List[Query] = []
         self._search_field = None
 
         self.operator = operator
         self.value = value
-        self.distance = distance
         if isinstance(search_field, str):
             self.search_field = SearchField(search_field)
         else:
@@ -198,24 +193,6 @@ def operator(self, is_op: bool) -> None:
             raise TypeError("operator must be a boolean")
         self._operator = is_op
 
-    @property
-    def distance(self) -> typing.Optional[int]:
-        """Distance property."""
-        return self._distance
-
-    @distance.setter
-    def distance(self, dist: typing.Optional[int]) -> None:
-        """Set distance property."""
-
-        if self.operator and self.value in {Operators.NEAR, Operators.WITHIN}:
-            if dist is None:
-                raise ValueError(f"{self.value} operator requires a distance")
-        else:
-            if dist is not None:
-                raise ValueError(f"{self.value} operator cannot have a distance")
-
-        self._distance = dist
-
     @property
     def children(self) -> typing.List[Query]:
         """Children property."""
@@ -284,38 +261,13 @@ def selects(self, *, record_dict: dict) -> bool:
         QueryTranslator.move_fields_to_terms(query_with_term_fields)
 
         # pylint: disable=protected-access
-        return query_with_term_fields._selects(record_dict=record_dict)
-
-    def _selects(self, record_dict: dict) -> bool:
-        if self.value == Operators.NOT:
-            return not self.children[0].selects(record_dict=record_dict)
-
-        if self.value == Operators.AND:
-            return all(x.selects(record_dict=record_dict) for x in self.children)
+        return query_with_term_fields.selects_record(record_dict=record_dict)
 
-        if self.value == Operators.OR:
-            return any(x.selects(record_dict=record_dict) for x in self.children)
-
-        assert not self.operator
-
-        assert self.search_field is not None, "Search field must be set for terms"
-        if self.search_field.value == Fields.TITLE:
-            field_value = record_dict.get("title", "").lower()
-        elif self.search_field.value == Fields.ABSTRACT:
-            field_value = record_dict.get("abstract", "").lower()
-        else:
-            raise ValueError(f"Unsupported search field: {self.search_field}")
-
-        value = self.value.lower().lstrip('"').rstrip('"')
-
-        # Handle wildcards
-        if "*" in value:
-            pattern = re.compile(value.replace("*", ".*").lower())
-            match = pattern.search(field_value)
-            return match is not None
-
-        # Match exact word
-        return value.lower() in field_value
+    def selects_record(self, record_dict: dict) -> bool:
+        """Abstract method to be implemented by subclasses to select records."""
+        raise NotImplementedError(
+            "This method should be implemented by subclasses of Query"
+        )
 
     def _get_confusion_matrix(self, records_dict: dict) -> dict:
         relevant_ids = set()

search_query/query_and.py

@@ -57,3 +57,6 @@ def children(self, children: typing.List[Query]) -> None:
         # Add each new child using add_child (ensures parent is set)
         for child in children or []:
             self.add_child(child)
+
+    def selects_record(self, record_dict: dict) -> bool:
+        return all(x.selects(record_dict=record_dict) for x in self.children)

search_query/query_near.py

@@ -1,9 +1,14 @@
 #!/usr/bin/env python
 """NEAR Query"""
 import typing
+from typing import cast
+from typing import List
+from typing import Union
 
+from search_query.constants import Operators
 from search_query.query import Query
 from search_query.query import SearchField
+from search_query.query_term import Term
 
 
 class NEARQuery(Query):
@@ -19,7 +24,7 @@ def __init__(
         *,
         search_field: typing.Optional[typing.Union[SearchField, str]] = None,
         position: typing.Optional[typing.Tuple[int, int]] = None,
-        distance: typing.Optional[int] = None,
+        distance: int,
         platform: str = "generic",
     ) -> None:
         """init method
@@ -29,18 +34,41 @@ def __init__(
         search field: search field to which the query should be applied
         """
 
+        query_children = [
+            c if isinstance(c, Query) else Term(value=c) for c in children
+        ]
+
         super().__init__(
             value=value,
-            children=children,
+            children=cast(List[Union[str, Query]], query_children),
             search_field=search_field
             if isinstance(search_field, SearchField)
             else SearchField(search_field)
             if search_field is not None
             else None,
             position=position,
-            distance=distance,
             platform=platform,
         )
+        self.children = query_children
+        self.distance: int = distance
+
+    @property
+    def distance(self) -> typing.Optional[int]:
+        """Distance property."""
+        return self._distance
+
+    @distance.setter
+    def distance(self, dist: typing.Optional[int]) -> None:
+        """Set distance property."""
+
+        if self.operator and self.value in {Operators.NEAR, Operators.WITHIN}:
+            if dist is None:
+                raise ValueError(f"{self.value} operator requires a distance")
+        else:
+            if dist is not None:
+                raise ValueError(f"{self.value} operator cannot have a distance")
+
+        self._distance = dist
 
     @property
     def children(self) -> typing.List[Query]:
@@ -51,13 +79,55 @@ def children(self) -> typing.List[Query]:
     def children(self, children: typing.List[Query]) -> None:
         """Set the children of NEAR query, updating parent pointers."""
         # Clear existing children and reset parent links (if necessary)
+
         self._children.clear()
+
         if not isinstance(children, list):
             raise TypeError("children must be a list of Query instances or strings")
 
-        if len(children) != 2:
-            raise ValueError("A NEAR query must have two children")
+        if self.platform != "deactivated":  # Note: temporary for EBSCO parser
+            if len(children) != 2:
+                raise ValueError("A NEAR query must have two children")
 
         # Add each new child using add_child (ensures parent is set)
         for child in children or []:
             self.add_child(child)
+
+    def selects_record(self, record_dict: dict) -> bool:
+        """Check if the record matches the NEAR query."""
+        assert len(self.children) == 2, "NEAR query must have two children"
+        assert self.children[0].search_field, "First child must have a search field"
+        assert self.children[1].search_field, "Second child must have a search field"
+        assert self.distance is not None, "NEAR query must have a distance"
+        assert (
+            self.children[0].search_field.value == self.children[1].search_field.value
+        ), "Both children of NEAR query must have the same search field"
+
+        # the self.children[0].value
+        # must be in self.distance words of self.children[1].value
+        field = self.children[0].search_field.value
+        text = record_dict.get(field, "")
+        if not isinstance(text, str):
+            return False
+
+        term1 = self.children[0].value.lower()
+        term2 = self.children[1].value.lower()
+
+        tokens = (
+            text.split()
+        )  # Simple whitespace tokenizer; can be replaced with a smarter one
+        # Get all positions of term1 and term2
+        positions_term1 = [
+            i for i, token in enumerate(tokens) if token.lower() == term1
+        ]
+        positions_term2 = [
+            i for i, token in enumerate(tokens) if token.lower() == term2
+        ]
+
+        # Check if any pair is within the allowed distance
+        for p1 in positions_term1:
+            for p2 in positions_term2:
+                if abs(p1 - p2) <= self.distance:
+                    return True
+
+        return False

search_query/query_not.py

@@ -57,3 +57,8 @@ def children(self, children: typing.List[Query]) -> None:
         # Add each new child using add_child (ensures parent is set)
         for child in children or []:
             self.add_child(child)
+
+    def selects_record(self, record_dict: dict) -> bool:
+        return self.children[0].selects(record_dict=record_dict) and not self.children[
+            1
+        ].selects(record_dict=record_dict)

search_query/query_or.py

@@ -58,3 +58,6 @@ def children(self, children: typing.List[Query]) -> None:
         # Add each new child using add_child (ensures parent is set)
         for child in children or []:
             self.add_child(child)
+
+    def selects_record(self, record_dict: dict) -> bool:
+        return any(x.selects(record_dict=record_dict) for x in self.children)

search_query/query_range.py

@@ -57,3 +57,28 @@ def children(self, children: typing.List[Query]) -> None:
         # Add each new child using add_child (ensures parent is set)
         for child in children or []:
             self.add_child(child)
+
+    def selects_record(self, record_dict: dict) -> bool:
+        """Check if the record matches the range query."""
+        assert len(self.children) == 2, "RANGE query must have two children"
+        assert self.children[0].search_field, "First child must have a search field"
+        assert self.children[1].search_field, "Second child must have a search field"
+        assert (
+            self.children[0].search_field.value == self.children[1].search_field.value
+        ), "Both children of RANGE query must have the same search field"
+
+        term1 = self.children[0].value.lower()
+        term2 = self.children[1].value.lower()
+        record_field = record_dict.get(
+            self.children[0].search_field.value, record_dict.get("year", "")
+        )
+
+        if term1.isdigit() and term2.isdigit() and record_field.isdigit():
+            value1 = int(term1)
+            value2 = int(term2)
+            record_value = int(record_field)
+            return value1 <= record_value <= value2
+
+        # Match other cases here (e.g., dates)
+
+        raise ValueError("Both children of RANGE query must be numeric values")

search_query/query_term.py

@@ -2,8 +2,10 @@
 """Query class."""
 from __future__ import annotations
 
+import re
 import typing
 
+from search_query.constants import Fields
 from search_query.constants import SearchField
 from search_query.query import Query
 
@@ -27,3 +29,23 @@ def __init__(
             position=position,
             platform=platform,
         )
+
+    def selects_record(self, record_dict: dict) -> bool:
+        assert self.search_field is not None, "Search field must be set for terms"
+        if self.search_field.value == Fields.TITLE:
+            field_value = record_dict.get("title", "").lower()
+        elif self.search_field.value == Fields.ABSTRACT:
+            field_value = record_dict.get("abstract", "").lower()
+        else:
+            raise ValueError(f"Unsupported search field: {self.search_field}")
+
+        value = self.value.lower().lstrip('"').rstrip('"')
+
+        # Handle wildcards
+        if "*" in value:
+            pattern = re.compile(value.replace("*", ".*").lower())
+            match = pattern.search(field_value)
+            return match is not None
+
+        # Match exact word
+        return value.lower() in field_value

search_query/wos/parser.py

@@ -16,6 +16,7 @@
 from search_query.parser_base import QueryStringParser
 from search_query.query import Query
 from search_query.query import SearchField
+from search_query.query_near import NEARQuery
 from search_query.query_term import Term
 from search_query.wos.constants import search_field_general_to_syntax
 from search_query.wos.linter import WOSQueryListLinter
@@ -241,12 +242,26 @@ def _handle_closing_parenthesis(
 
         # Return the operator and children if there is an operator
         if current_operator:
+            if distance:
+                # If there is a distance, it must be a proximity operator
+                if current_operator not in {"NEAR", "WITHIN"}:
+                    raise ValueError(
+                        f"Distance {distance} "
+                        "is only allowed for NEAR or WITHIN operators, "
+                        f"not {current_operator}"
+                    )
+                return NEARQuery(
+                    value=current_operator,
+                    children=children,
+                    search_field=search_field,
+                    platform="deactivated",
+                    distance=distance,
+                )
             return Query(
                 value=current_operator,
                 children=children,
                 search_field=search_field,
                 platform="deactivated",
-                distance=distance,
             )
 
         # Multiple children without operator are not allowed