insights/common/scriptvm/python/utils.py

from typing import Any

import re2

from common.scriptvm.python.objects import is_script_date, is_script_datetime
from common.scriptvm.python.stl.date import date_string_to_seconds, to_script_datetime

_CASE_INSENSITIVE_OPTS = re2.Options()
_CASE_INSENSITIVE_OPTS.case_sensitive = False

COST_PER_UNIT = 8


def _temporal_seconds(value: Any) -> float | None:
    """Epoch seconds for a Script datetime/date value, else None. A bare ScriptDate is UTC midnight."""
    if is_script_datetime(value):
        return value["dt"]
    if is_script_date(value):
        return to_script_datetime(value)["dt"]
    return None


class ScriptVMException(Exception):
    pass


class UncaughtScriptVMException(ScriptVMException):
    type: str
    message: str
    payload: Any

    def __init__(self, type, message, payload):
        super().__init__(message)
        self.type = type
        self.message = message
        self.payload = payload

    def __str__(self):
        msg = self.message.replace("'", "\\'")
        return f"{self.type}('{msg}')"


class ScriptVMRuntimeExceededException(ScriptVMException):
    """Exception thrown when ScriptVM code exceeds its runtime limit"""

    def __init__(self, timeout_seconds: float, ops_performed: int):
        self.timeout_seconds = timeout_seconds
        self.ops_performed = ops_performed
        super().__init__(f"Runtime exceeded {timeout_seconds} seconds after {ops_performed} operations")


class ScriptVMMemoryExceededException(ScriptVMException):
    """Exception thrown when ScriptVM code exceeds its memory limit"""

    def __init__(self, memory_limit: int, attempted_memory: int):
        self.memory_limit = memory_limit
        self.attempted_memory = attempted_memory
        super().__init__(f"Memory limit of {memory_limit} bytes exceeded. Attempted to use {attempted_memory} bytes")


def _require_string(value: Any, name: str, function_name: str) -> str:
    if not isinstance(value, str):
        raise ScriptVMException(f"Function {function_name} requires {name} to be a string, got {type(value).__name__}")
    return value


def _format_regex_error(error: Exception) -> str:
    if error.args and isinstance(error.args[0], bytes):
        return error.args[0].decode("utf-8", errors="replace")
    return str(error)


def _compile_regex(pattern: str, case_insensitive: bool = False) -> Any:
    try:
        return re2.compile(pattern, options=_CASE_INSENSITIVE_OPTS) if case_insensitive else re2.compile(pattern)
    except re2.error as e:
        raise ScriptVMException(f"Invalid regex pattern: {_format_regex_error(e)}") from e


def regex_match(string: Any, pattern: Any, case_insensitive: bool = False) -> bool:
    if not string or not pattern:
        return False

    string = _require_string(string, "input", "match")
    pattern = _require_string(pattern, "pattern", "match")
    return _compile_regex(pattern, case_insensitive).search(string) is not None


def like(string: Any, pattern: Any, case_insensitive: bool = False) -> bool:
    pattern = re2.escape(pattern).replace("%", ".*").replace("_", ".")
    re_pattern = re2.compile(pattern, options=_CASE_INSENSITIVE_OPTS) if case_insensitive else re2.compile(pattern)
    return re_pattern.search(string) is not None


def get_nested_value(obj, chain, nullish=False) -> Any:
    if obj is None:
        return None
    for key in chain:
        if nullish and obj is None:
            return None
        if isinstance(key, int):
            if key == 0:
                raise ScriptVMException(f"Script arrays start from index 1")
            elif key > 0:
                if key > len(obj):
                    return None
                obj = obj[key - 1]
            elif key < 0:
                if -key > len(obj):
                    return None
                obj = obj[key]
        else:
            obj = obj.get(key, None)
    return obj


def set_nested_value(obj, chain, value) -> Any:
    if obj is None:
        return None
    for key in chain[:-1]:
        if isinstance(key, int):
            obj = obj[key]
        else:
            obj = obj.get(key, None)

    if isinstance(obj, dict):
        obj[chain[-1]] = value
    elif isinstance(obj, list):
        if not isinstance(chain[-1], int):
            raise ScriptVMException(f"Invalid index: {chain[-1]}")
        if chain[-1] <= 0:
            raise ScriptVMException(f"Script arrays start from index 1")
        if chain[-1] > len(obj):
            raise ScriptVMException(f"Index {chain[-1]} out of range for array of length {len(obj)}")
        obj[chain[-1] - 1] = value
    else:
        raise ScriptVMException(f'Can not set property "{chain[-1]}" on object of type "{type(obj).__name__}"')

    return obj


def calculate_cost(object, marked: set | None = None) -> int:
    if marked is None:
        marked = set()
    if isinstance(object, dict) or isinstance(object, list) or isinstance(object, tuple):
        if id(object) in marked:
            return COST_PER_UNIT
        marked.add(id(object))
        try:
            if isinstance(object, dict):
                return COST_PER_UNIT + sum(
                    [calculate_cost(key, marked) + calculate_cost(value, marked) for key, value in object.items()]
                )
            elif isinstance(object, list) or isinstance(object, tuple):
                return COST_PER_UNIT + sum([calculate_cost(val, marked) for val in object])
        finally:
            marked.remove(id(object))
    elif isinstance(object, str):
        return COST_PER_UNIT + len(object)
    return COST_PER_UNIT


def unify_comparison_types(left, right):
    # Two temporal values order by epoch seconds (matching Datastore and the TS/Rust VMs). Without
    # this a ScriptDateTime/ScriptDate dict falls through unchanged and ordering operators end up comparing
    # dicts, which Python can't order.
    left_seconds = _temporal_seconds(left)
    right_seconds = _temporal_seconds(right)
    if left_seconds is not None and right_seconds is not None:
        return left_seconds, right_seconds
    # A bare-field SQL comparison like `timestamp > toDateTime(...)` puts a plain date-like string
    # against a ScriptDateTime/ScriptDate object. Parse the string the same way `toDateTime` would rather
    # than falling through to the string/number branches below, which leave it unordered.
    if left_seconds is not None and isinstance(right, str):
        right_seconds_from_string = date_string_to_seconds(right)
        if right_seconds_from_string is not None:
            return left_seconds, right_seconds_from_string
    if right_seconds is not None and isinstance(left, str):
        left_seconds_from_string = date_string_to_seconds(left)
        if left_seconds_from_string is not None:
            return left_seconds_from_string, right_seconds

    # Handle boolean cases FIRST since bool is a subclass of int in Python
    if isinstance(left, bool) and isinstance(right, str):
        # Convert string to boolean: 'true'/'false' strings, or truthy/falsy
        if right.lower() == "true":
            return left, True
        elif right.lower() == "false":
            return left, False
        else:
            return left, bool(right)
    if isinstance(left, str) and isinstance(right, bool):
        # Convert string to boolean: 'true'/'false' strings, or truthy/falsy
        if left.lower() == "true":
            return True, right
        elif left.lower() == "false":
            return False, right
        else:
            return bool(left), right
    if isinstance(left, bool) and isinstance(right, int | float):
        return int(left), right
    if isinstance(left, int | float) and isinstance(right, bool):
        return left, int(right)

    # Handle numeric conversions (after boolean checks)
    if isinstance(left, int | float) and isinstance(right, str):
        try:
            return left, float(right)
        except ValueError:
            return left, right
    if isinstance(left, str) and isinstance(right, int | float):
        try:
            return float(left), right
        except ValueError:
            return left, right

    return left, right