"""Array and array-transform Python evaluators."""
from __future__ import annotations
from datetime import datetime
from typing import Any
from .base import BasePythonMixin
# BSON comparison order groups (lower = earlier in sort).
# Order: Null < Numbers < Strings < Objects < Arrays <
# BinData < ObjectId < Boolean < Date < Timestamp < Regex
[docs]
def _bson_sort_key(value: Any) -> tuple[int, Any]:
"""Return a sort key tuple that encodes BSON comparison order.
The first element is the BSON type group (0..9), the second is the
value itself for same-type comparison. Unknown types fall at the end.
"""
if value is None:
return (0, None)
# bool is a subclass of int; check it first.
if isinstance(value, bool):
return (6, value)
if isinstance(value, (int, float)):
return (1, value)
if isinstance(value, str):
return (2, value)
if isinstance(value, dict):
return (3, str(value))
if isinstance(value, list):
return (4, str(value))
if isinstance(value, (bytes, bytearray)):
return (5, value)
if isinstance(value, datetime):
return (7, value)
# From neosqlite.objectid / neosqlite.binary if available
type_name = type(value).__name__
if type_name == "ObjectId":
return (5, str(value))
if type_name == "Binary":
return (5, bytes(value) if hasattr(value, "__bytes__") else str(value))
# Fallback: unknown types sort last
return (99, str(value))
[docs]
def _bson_sort(array: list[Any]) -> list[Any]:
"""Sort a list using MongoDB BSON comparison order."""
try:
return sorted(array, key=_bson_sort_key)
except TypeError:
# Fallback: if two values of the same type still can't be compared,
# return original order.
return array
[docs]
def _set_key(v: Any) -> Any:
"""Hashable canonical form of an element for $set* operators (#122)."""
import json as _json
if isinstance(v, dict):
return ("d", _json.dumps(v, sort_keys=True, default=str))
if isinstance(v, list):
return ("l", tuple(_set_key(x) for x in v))
if isinstance(v, bool):
return ("b", v)
if isinstance(v, int):
return ("i", v)
if isinstance(v, float):
return ("f", repr(v))
if isinstance(v, str):
return ("s", v)
return ("o", str(v))
[docs]
def _dedupe(seq: list) -> list:
seen: set = set()
out: list = []
for x in seq:
k = _set_key(x)
if k not in seen:
seen.add(k)
out.append(x)
return out
[docs]
def _bson_min_max(array: list, is_min: bool):
"""min/max under BSON type ordering; None if no non-null values."""
from neosqlite.collection.type_utils import bson_sort_key
vals = [v for v in array if v is not None]
if not vals:
return None
key = lambda v: bson_sort_key(v) # noqa: E731
return min(vals, key=key) if is_min else max(vals, key=key)
[docs]
class ArrayPythonMixin(BasePythonMixin):
"""Array, set, and transform ($filter/$map/$reduce) operators."""
[docs]
def _evaluate_array_python(
self, operator: str, operands: list[Any], document: dict[str, Any]
) -> Any:
"""Evaluate array operators in Python."""
# Normalize operands for operators that accept single values
if operator in (
"$size",
"$isArray",
"$sum",
"$avg",
"$min",
"$max",
) and not isinstance(operands, list):
operands = [operands]
match operator:
case "$size":
if len(operands) != 1:
raise ValueError("$size requires exactly 1 operand")
array = self._evaluate_operand_python(operands[0], document)
if isinstance(array, list):
return len(array)
return None
case "$in":
if len(operands) != 2:
raise ValueError("$in requires exactly 2 operands")
value = self._evaluate_operand_python(operands[0], document)
array = self._evaluate_operand_python(operands[1], document)
if isinstance(array, list):
return value in array
return False
case "$isArray":
if len(operands) != 1:
raise ValueError("$isArray requires exactly 1 operand")
value = self._evaluate_operand_python(operands[0], document)
return isinstance(value, list)
case "$sum" | "$avg" | "$min" | "$max":
# Handle both list and single operand formats
if not isinstance(operands, list):
array_ops = [operands]
else:
array_ops = operands
if len(array_ops) != 1:
raise ValueError(f"{operator} requires exactly 1 operand")
array = self._evaluate_operand_python(array_ops[0], document)
if not isinstance(array, list):
return 0 if operator == "$sum" else None
# Filter numeric values for sum/avg
nums = [
v
for v in array
if isinstance(v, (int, float)) and not isinstance(v, bool)
]
if not nums:
if operator == "$sum":
return 0
return None
match operator:
case "$sum":
return sum(nums)
case "$avg":
return sum(nums) / len(nums)
case "$min":
# BSON-ordered min: mixed-type arrays must not
# raise TypeError (#122)
return _bson_min_max(array, True)
case "$max":
return _bson_min_max(array, False)
case _:
return None
case "$arrayElemAt":
if len(operands) != 2:
raise ValueError("$arrayElemAt requires exactly 2 operands")
array = self._evaluate_operand_python(operands[0], document)
index = self._evaluate_operand_python(operands[1], document)
if isinstance(array, list) and isinstance(index, int):
try:
return array[index]
except IndexError:
return None
return None
case "$first":
# Handle both list and single operand formats
if not isinstance(operands, list):
ops = [operands]
else:
ops = operands
if len(ops) != 1:
raise ValueError("$first requires exactly 1 operand")
array = self._evaluate_operand_python(ops[0], document)
if isinstance(array, list) and array:
return array[0]
return None
case "$last":
# Handle both list and single operand formats
if not isinstance(operands, list):
ops = [operands]
else:
ops = operands
if len(ops) != 1:
raise ValueError("$last requires exactly 1 operand")
array = self._evaluate_operand_python(ops[0], document)
if isinstance(array, list) and array:
return array[-1]
return None
case "$firstN":
# Get first N elements from array
# MongoDB syntax: { $firstN: { input: <array>, n: <number> } }
if isinstance(operands, dict):
array_operand = operands.get("input")
n_operand = operands.get("n")
elif isinstance(operands, list) and len(operands) == 2:
array_operand = operands[0]
n_operand = operands[1]
else:
raise ValueError("$firstN requires input array and n count")
array = self._evaluate_operand_python(array_operand, document)
n = self._evaluate_operand_python(n_operand, document)
if not isinstance(array, list) or n is None:
return []
if not isinstance(n, (int, float)) or isinstance(n, bool):
raise ValueError("$firstN requires a numeric n")
if n <= 0:
# MongoDB errors on non-positive n; $lastN already does
raise ValueError("$firstN requires a positive n")
return array[: int(n)]
case "$lastN":
# Get last N elements from array
# MongoDB syntax: { $lastN: { input: <array>, n: <number> } }
if isinstance(operands, dict):
array_operand = operands.get("input")
n_operand = operands.get("n")
elif isinstance(operands, list) and len(operands) == 2:
array_operand = operands[0]
n_operand = operands[1]
else:
raise ValueError("$lastN requires input array and n count")
array = self._evaluate_operand_python(array_operand, document)
n = self._evaluate_operand_python(n_operand, document)
if not isinstance(array, list) or n is None:
return []
n_int = int(n)
if n_int <= 0:
return []
return array[-n_int:] if n_int < len(array) else array
case "$maxN":
# Get maximum N elements from array (sorted descending, take first N)
# MongoDB syntax: { $maxN: { input: <array>, n: <number> } }
if isinstance(operands, dict):
array_operand = operands.get("input")
n_operand = operands.get("n")
elif isinstance(operands, list) and len(operands) == 2:
array_operand = operands[0]
n_operand = operands[1]
else:
raise ValueError("$maxN requires input array and n count")
array = self._evaluate_operand_python(array_operand, document)
n = self._evaluate_operand_python(n_operand, document)
if not isinstance(array, list) or n is None:
return []
# Sort descending and take first N
try:
sorted_array = sorted(array, reverse=True)
return sorted_array[: int(n)]
except (TypeError, ValueError):
return []
case "$minN":
# Get minimum N elements from array (sorted ascending, take first N)
# MongoDB syntax: { $minN: { input: <array>, n: <number> } }
if isinstance(operands, dict):
array_operand = operands.get("input")
n_operand = operands.get("n")
elif isinstance(operands, list) and len(operands) == 2:
array_operand = operands[0]
n_operand = operands[1]
else:
raise ValueError("$minN requires input array and n count")
array = self._evaluate_operand_python(array_operand, document)
n = self._evaluate_operand_python(n_operand, document)
if not isinstance(array, list) or n is None:
return []
# Sort ascending and take first N
try:
sorted_array = sorted(array)
return sorted_array[: int(n)]
except (TypeError, ValueError):
return []
case "$sortArray":
# Sort array elements
# MongoDB syntax: { $sortArray: { input: <array>, sortBy: { <field>: <direction> } } }
if isinstance(operands, dict):
array_operand = operands.get("input")
sort_by = operands.get("sortBy")
elif isinstance(operands, list) and len(operands) >= 1:
array_operand = operands[0]
sort_by = operands[1] if len(operands) > 1 else None
else:
raise ValueError("$sortArray requires input array")
array = self._evaluate_operand_python(array_operand, document)
if not isinstance(array, list):
return []
# If no sortBy specified, sort by value using BSON comparison order.
# BSON order: Null < Numbers < Strings < Objects < Arrays <
# BinData < ObjectId < Boolean < Date < Timestamp < Regex
if sort_by is None:
return _bson_sort(array)
# Numeric sortBy: 1 (ascending) or -1 (descending) by value.
if (
isinstance(sort_by, (int, float))
and not isinstance(sort_by, bool)
and sort_by in (1, -1)
):
sorted_array = _bson_sort(array)
if sort_by == -1:
return list(reversed(sorted_array))
return sorted_array
# Sort by field (for array of objects)
if isinstance(sort_by, dict):
# Get first field and direction
sort_field = next(iter(sort_by.keys()))
direction = sort_by[sort_field]
reverse = direction == -1
try:
def sort_key(x: Any) -> Any:
"""
Extract the sort field from a dictionary or return the value.
"""
return (
x.get(sort_field) if isinstance(x, dict) else x
)
return sorted(
array,
key=sort_key, # type: ignore[arg-type]
reverse=reverse,
)
except (TypeError, AttributeError):
return array
return array
case "$slice":
if not isinstance(operands, list) or len(operands) < 2:
raise ValueError("$slice requires array and count/position")
array = self._evaluate_operand_python(operands[0], document)
count = self._evaluate_operand_python(operands[1], document)
if not isinstance(array, list):
return []
if len(operands) >= 3:
skip = self._evaluate_operand_python(operands[2], document)
return array[skip : skip + count]
elif isinstance(count, int) and count < 0:
return array[count:]
else:
return array[:count]
case "$indexOfArray":
if len(operands) != 2:
raise ValueError(
"$indexOfArray requires exactly 2 operands"
)
array = self._evaluate_operand_python(operands[0], document)
value = self._evaluate_operand_python(operands[1], document)
if isinstance(array, list):
try:
return array.index(value)
except ValueError:
return -1
return -1
case "$setEquals":
if len(operands) != 2:
raise ValueError("$setEquals requires exactly 2 operands")
set1 = self._evaluate_operand_python(operands[0], document)
set2 = self._evaluate_operand_python(operands[1], document)
if isinstance(set1, list) and isinstance(set2, list):
return set(set1) == set(set2)
return False
case "$setIntersection":
if len(operands) != 2:
raise ValueError(
"$setIntersection requires exactly 2 operands"
)
set1 = self._evaluate_operand_python(operands[0], document)
set2 = self._evaluate_operand_python(operands[1], document)
if isinstance(set1, list) and isinstance(set2, list):
b_keys = {_set_key(y) for y in set2}
return _dedupe([x for x in set1 if _set_key(x) in b_keys])
return []
case "$setUnion":
if len(operands) != 2:
raise ValueError("$setUnion requires exactly 2 operands")
set1 = self._evaluate_operand_python(operands[0], document)
set2 = self._evaluate_operand_python(operands[1], document)
if isinstance(set1, list) and isinstance(set2, list):
return _dedupe(list(set1) + list(set2))
return []
case "$setDifference":
if len(operands) != 2:
raise ValueError(
"$setDifference requires exactly 2 operands"
)
set1 = self._evaluate_operand_python(operands[0], document)
set2 = self._evaluate_operand_python(operands[1], document)
if isinstance(set1, list) and isinstance(set2, list):
b_keys = {_set_key(y) for y in set2}
return [x for x in set1 if _set_key(x) not in b_keys]
return []
case "$setIsSubset":
if len(operands) != 2:
raise ValueError("$setIsSubset requires exactly 2 operands")
set1 = self._evaluate_operand_python(operands[0], document)
set2 = self._evaluate_operand_python(operands[1], document)
if isinstance(set1, list) and isinstance(set2, list):
return set(set1).issubset(set(set2))
return False
case "$anyElementTrue":
# Handle both list and single operand formats
if not isinstance(operands, list):
operands = [operands]
if len(operands) != 1:
raise ValueError(
"$anyElementTrue requires exactly 1 operand"
)
array = self._evaluate_operand_python(operands[0], document)
if isinstance(array, list):
return any(array)
return False
case "$allElementsTrue":
# Handle both list and single operand formats
if not isinstance(operands, list):
operands = [operands]
if len(operands) != 1:
raise ValueError(
"$allElementsTrue requires exactly 1 operand"
)
array = self._evaluate_operand_python(operands[0], document)
if isinstance(array, list):
return all(array)
return False
case _:
raise NotImplementedError(
f"Array operator {operator} not supported in Python evaluation"
)