query.py

"""
クエリパーサー - MongoDBライクなクエリ構文をサポート
"""

import re
from typing import Any, Dict, List, Union
from datetime import datetime


class Query:
    """MongoDBライクなクエリを処理するクラス"""
    
    def __init__(self, query_dict: Dict[str, Any]):
        self.query = query_dict
    
    def matches(self, document: Dict[str, Any]) -> bool:
        """
        ドキュメントがクエリにマッチするかを判定
        
        Args:
            document: 検索対象のドキュメント
            
        Returns:
            bool: マッチする場合はTrue
        """
        return self._evaluate_query(self.query, document)
    
    def _evaluate_query(self, query: Dict[str, Any], document: Dict[str, Any]) -> bool:
        """
        クエリを評価してドキュメントがマッチするかを判定
        
        Args:
            query: 評価するクエリ
            document: 検索対象のドキュメント
            
        Returns:
            bool: マッチする場合はTrue
        """
        if not query:
            return True
            
        for key, value in query.items():
            if key.startswith('$'):
                # 論理オペレーター
                if not self._evaluate_logical_operator(key, value, document):
                    return False
            else:
                # フィールド比較
                if not self._evaluate_field_query(key, value, document):
                    return False
        
        return True
    
    def _evaluate_logical_operator(self, operator: str, value: Any, document: Dict[str, Any]) -> bool:
        """
        論理オペレーターを評価
        
        Args:
            operator: 論理オペレーター ($and, $or, $not, $nor)
            value: オペレーターの値
            document: 検索対象のドキュメント
            
        Returns:
            bool: 評価結果
        """
        if operator == '$and':
            return all(self._evaluate_query(q, document) for q in value)
        elif operator == '$or':
            return any(self._evaluate_query(q, document) for q in value)
        elif operator == '$not':
            return not self._evaluate_query(value, document)
        elif operator == '$nor':
            return not any(self._evaluate_query(q, document) for q in value)
        else:
            return False
    
    def _evaluate_field_query(self, field: str, query_value: Any, document: Dict[str, Any]) -> bool:
        """
        フィールドクエリを評価
        
        Args:
            field: フィールド名
            query_value: クエリの値
            document: 検索対象のドキュメント
            
        Returns:
            bool: 評価結果
        """
        # ネストしたフィールドの取得
        doc_value = self._get_nested_field(document, field)
        
        if isinstance(query_value, dict):
            # 比較オペレーターの処理
            return self._evaluate_comparison_operators(doc_value, query_value)
        else:
            # 単純な等価比較
            return doc_value == query_value
    
    def _get_nested_field(self, document: Dict[str, Any], field: str) -> Any:
        """
        ネストしたフィールドの値を取得
        
        Args:
            document: 検索対象のドキュメント
            field: フィールド名（ドット記法サポート）
            
        Returns:
            Any: フィールドの値
        """
        keys = field.split('.')
        value = document
        
        for key in keys:
            if isinstance(value, dict) and key in value:
                value = value[key]
            else:
                return None
        
        return value
    
    def _evaluate_comparison_operators(self, doc_value: Any, query_ops: Dict[str, Any]) -> bool:
        """
        比較オペレーターを評価
        
        Args:
            doc_value: ドキュメントの値
            query_ops: クエリオペレーターのDict
            
        Returns:
            bool: 評価結果
        """
        for op, op_value in query_ops.items():
            if op == '$eq':
                if doc_value != op_value:
                    return False
            elif op == '$ne':
                if doc_value == op_value:
                    return False
            elif op == '$gt':
                if doc_value is None or doc_value <= op_value:
                    return False
            elif op == '$gte':
                if doc_value is None or doc_value < op_value:
                    return False
            elif op == '$lt':
                if doc_value is None or doc_value >= op_value:
                    return False
            elif op == '$lte':
                if doc_value is None or doc_value > op_value:
                    return False
            elif op == '$in':
                if not isinstance(op_value, list):
                    return False
                if isinstance(doc_value, list):
                    # 配列フィールドの場合、いずれかの要素がマッチするか
                    if not any(item in op_value for item in doc_value):
                        return False
                else:
                    # 単一値フィールドの場合
                    if doc_value not in op_value:
                        return False
            elif op == '$nin':
                if not isinstance(op_value, list):
                    return False
                if isinstance(doc_value, list):
                    # 配列フィールドの場合、すべての要素がマッチしないか
                    if any(item in op_value for item in doc_value):
                        return False
                else:
                    # 単一値フィールドの場合
                    if doc_value in op_value:
                        return False
            elif op == '$regex':
                if not self._evaluate_regex(doc_value, op_value, query_ops.get('$options')):
                    return False
            elif op == '$exists':
                exists = doc_value is not None
                if exists != op_value:
                    return False
            elif op == '$type':
                if not self._evaluate_type(doc_value, op_value):
                    return False
            elif op == '$size':
                if not isinstance(doc_value, (list, str)) or len(doc_value) != op_value:
                    return False
            elif op == '$all':
                if not isinstance(doc_value, list) or not all(item in doc_value for item in op_value):
                    return False
            elif op == '$elemMatch':
                if not isinstance(doc_value, list):
                    return False
                
                # 配列内の各要素をチェック
                for item in doc_value:
                    if isinstance(item, dict):
                        # オブジェクトの場合、通常のクエリ評価
                        if self._evaluate_query(op_value, item):
                            break  # 一つでもマッチしたらOK
                    else:
                        # プリミティブ値の場合、一時的なドキュメントを作成
                        for op_key, op_val in op_value.items():
                            if op_key.startswith('$'):
                                # 比較オペレーターを直接適用
                                if self._evaluate_comparison_operators(item, {op_key: op_val}):
                                    break
                        else:
                            continue
                        break
                else:
                    return False
            elif op == '$mod':
                if not isinstance(doc_value, (int, float)) or doc_value % op_value[0] != op_value[1]:
                    return False
        
        return True
    
    def _evaluate_regex(self, doc_value: Any, pattern: Union[str, re.Pattern], options: str = None) -> bool:
        """
        正規表現の評価
        
        Args:
            doc_value: ドキュメントの値
            pattern: 正規表現パターン
            options: 正規表現オプション
            
        Returns:
            bool: マッチする場合はTrue
        """
        if not isinstance(doc_value, str):
            return False
        
        if isinstance(pattern, str):
            flags = 0
            if options:
                if 'i' in options:
                    flags |= re.IGNORECASE
                if 'm' in options:
                    flags |= re.MULTILINE
                if 's' in options:
                    flags |= re.DOTALL
                if 'x' in options:
                    flags |= re.VERBOSE
            
            try:
                return bool(re.search(pattern, doc_value, flags))
            except re.error:
                return False
        
        return bool(pattern.search(doc_value))
    
    def _evaluate_type(self, doc_value: Any, type_spec: Union[str, int]) -> bool:
        """
        型の評価
        
        Args:
            doc_value: ドキュメントの値
            type_spec: 型の指定
            
        Returns:
            bool: 型がマッチする場合はTrue
        """
        type_map = {
            'double': (float, int),
            'string': str,
            'object': dict,
            'array': list,
            'bool': bool,
            'null': type(None),
            'int': int,
            'date': datetime,
            # BSON type numbers
            1: (float, int),  # double
            2: str,           # string
            3: dict,          # object
            4: list,          # array
            8: bool,          # bool
            10: type(None),   # null
            16: int,          # int32
            18: int,          # int64
        }
        
        target_type = type_map.get(type_spec)
        if target_type is None:
            return False
        
        if isinstance(target_type, tuple):
            return isinstance(doc_value, target_type)
        else:
            return isinstance(doc_value, target_type)