krons 0.1.0__py3-none-any.whl

kronos/utils/fuzzy/_fuzzy_json.py

@@ -0,0 +1,288 @@
+# Copyright (c) 2025 - 2026, HaiyangLi <quantocean.li at gmail dot com>
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Any
+
+import orjson
+
+# Security limit: Maximum input string size for JSON parsing.
+# Large inputs can cause memory exhaustion during parsing and regex operations.
+# Default: 10MB - generous for most use cases while preventing DoS.
+MAX_JSON_INPUT_SIZE = 10 * 1024 * 1024  # 10 MB
+
+__all__ = ("fuzzy_json",)
+
+
+def fuzzy_json(
+    str_to_parse: str,
+    /,
+    *,
+    max_size: int = MAX_JSON_INPUT_SIZE,
+) -> dict[str, Any] | list[dict[str, Any]]:
+    """Parse JSON string with fuzzy error correction (quotes, spacing, brackets).
+
+    Security Note:
+        Input size is limited to prevent memory exhaustion attacks.
+        Uses orjson which has built-in recursion depth limits (CVE-2024-27454 fixed).
+
+    Limitation:
+        This is designed for "nearly-valid" JSON from LLMs (single quotes, trailing
+        commas, unquoted keys). It is NOT a security boundary for adversarial input.
+
+    Args:
+        str_to_parse: JSON string to parse
+        max_size: Maximum allowed input size in bytes (default: 10MB)
+
+    Returns:
+        Either a dict or a list of dicts. Will NOT return primitive types
+        (int, float, str, bool, None) or lists of primitives.
+
+    Raises:
+        TypeError: If input is not a string or parsed JSON is a primitive type.
+        ValueError: If input is empty, exceeds size limit, or parsing fails.
+    """
+    _check_valid_str(str_to_parse, max_size=max_size)
+
+    # Direct parse
+    try:
+        result = orjson.loads(str_to_parse)
+        return _validate_return_type(result)
+    except orjson.JSONDecodeError:
+        pass
+
+    # State-machine cleaning
+    cleaned = _clean_json_string_safe(str_to_parse)
+    try:
+        result = orjson.loads(cleaned)
+        return _validate_return_type(result)
+    except orjson.JSONDecodeError:
+        pass
+
+    # Bracket balancing
+    fixed = fix_json_string(cleaned)
+    try:
+        result = orjson.loads(fixed)
+        return _validate_return_type(result)
+    except orjson.JSONDecodeError:
+        pass
+
+    # If all attempts fail
+    raise ValueError("Invalid JSON string")
+
+
+def _check_valid_str(
+    str_to_parse: str,
+    /,
+    *,
+    max_size: int = MAX_JSON_INPUT_SIZE,
+) -> None:
+    """Validate input string for JSON parsing.
+
+    Args:
+        str_to_parse: Input string to validate
+        max_size: Maximum allowed size in bytes
+
+    Raises:
+        TypeError: If input is not a string
+        ValueError: If input is empty or exceeds size limit
+    """
+    if not isinstance(str_to_parse, str):
+        raise TypeError("Input must be a string")
+    if not str_to_parse.strip():
+        raise ValueError("Input string is empty")
+    if len(str_to_parse) > max_size:
+        msg = (
+            f"Input size ({len(str_to_parse)} bytes) exceeds maximum "
+            f"({max_size} bytes). This limit prevents memory exhaustion."
+        )
+        raise ValueError(msg)
+
+
+def _validate_return_type(result: Any) -> dict[str, Any] | list[dict[str, Any]]:
+    """Validate that parsed JSON matches the declared return type.
+
+    Args:
+        result: The result from orjson.loads()
+
+    Returns:
+        The validated result (dict or list[dict])
+
+    Raises:
+        TypeError: If result is a primitive type or list of non-dict elements
+    """
+    if isinstance(result, dict):
+        return result
+
+    if isinstance(result, list):
+        if not result:  # Empty list valid (vacuous truth)
+            return result
+        for i, item in enumerate(result):
+            if not isinstance(item, dict):
+                raise TypeError(
+                    f"fuzzy_json returns dict or list[dict], got list with "
+                    f"non-dict element at index {i}: {type(item).__name__}"
+                )
+        return result
+
+    raise TypeError(
+        f"fuzzy_json returns dict or list[dict], got primitive type: {type(result).__name__}"
+    )
+
+
+def _clean_json_string_safe(s: str) -> str:
+    """State-machine JSON cleaner that preserves string content.
+
+    Fixes common LLM output issues without corrupting quoted content:
+    - Single quotes -> double quotes (with escaping)
+    - Trailing commas before ] or }
+    - Unquoted object keys
+    """
+    result: list[str] = []
+    pos = 0
+    length = len(s)
+
+    while pos < length:
+        char = s[pos]
+
+        if char == "'":  # Convert single-quoted string to double-quoted
+            result.append('"')
+            pos += 1
+            while pos < length:
+                inner_char = s[pos]
+                if inner_char == "\\":  # Escape sequence
+                    if pos + 1 < length:
+                        next_char = s[pos + 1]
+                        if next_char == "'":  # \' -> apostrophe
+                            result.append("'")
+                            pos += 2
+                            continue
+                        result.append(inner_char)
+                        result.append(next_char)
+                        pos += 2
+                        continue
+                    result.append(inner_char)
+                    pos += 1
+                    continue
+                if inner_char == "'":  # End single-quoted string
+                    result.append('"')
+                    pos += 1
+                    break
+                if inner_char == '"':  # Escape inner double quotes
+                    result.append('\\"')
+                    pos += 1
+                    continue
+                result.append(inner_char)
+                pos += 1
+            continue
+
+        if char == '"':  # Pass through double-quoted strings
+            result.append(char)
+            pos += 1
+            while pos < length:
+                inner_char = s[pos]
+                if inner_char == "\\":  # Copy escape sequences
+                    result.append(inner_char)
+                    if pos + 1 < length:
+                        pos += 1
+                        result.append(s[pos])
+                    pos += 1
+                    continue
+                result.append(inner_char)
+                if inner_char == '"':
+                    pos += 1
+                    break
+                pos += 1
+            continue
+
+        if char in "{,":  # Handle unquoted keys and trailing commas
+            if char == ",":  # Check for trailing comma
+                lookahead = pos + 1
+                while lookahead < length and s[lookahead] in " \t\n\r":
+                    lookahead += 1
+                if lookahead < length and s[lookahead] in "]}":
+                    pos += 1  # Skip trailing comma
+                    continue
+
+            result.append(char)
+            pos += 1
+
+            while pos < length and s[pos] in " \t\n\r":  # Skip whitespace
+                result.append(s[pos])
+                pos += 1
+
+            if pos < length and s[pos] not in "\"'{[":  # Check for unquoted key
+                key_start = pos
+                while pos < length and (s[pos].isalnum() or s[pos] == "_"):
+                    pos += 1
+                if pos < length and key_start < pos:
+                    key_end = pos
+                    while pos < length and s[pos] in " \t\n\r":
+                        pos += 1
+                    if pos < length and s[pos] == ":":
+                        key = s[key_start:key_end]
+                        result.append(f'"{key}"')
+                        continue
+                    else:
+                        pos = key_start  # Not a key, restore
+            continue
+
+        result.append(char)
+        pos += 1
+
+    return "".join(result).strip()
+
+
+def fix_json_string(str_to_parse: str, /) -> str:
+    """Balance unmatched brackets in JSON string.
+
+    Args:
+        str_to_parse: JSON string with potentially missing closing brackets.
+
+    Returns:
+        JSON string with missing closing brackets appended.
+
+    Raises:
+        ValueError: If input is empty or has mismatched/extra brackets.
+    """
+    if not str_to_parse:
+        raise ValueError("Input string is empty")
+
+    brackets = {"{": "}", "[": "]"}
+    open_brackets = []
+    pos = 0
+    length = len(str_to_parse)
+
+    while pos < length:
+        char = str_to_parse[pos]
+
+        if char == "\\":
+            pos += 2
+            continue
+
+        if char == '"':
+            pos += 1
+            while pos < length:  # Skip string content
+                if str_to_parse[pos] == "\\":
+                    pos += 2
+                    continue
+                if str_to_parse[pos] == '"':
+                    pos += 1
+                    break
+                pos += 1
+            continue
+
+        if char in brackets:
+            open_brackets.append(brackets[char])
+        elif char in brackets.values():
+            if not open_brackets:
+                raise ValueError("Extra closing bracket found.")
+            if open_brackets[-1] != char:
+                raise ValueError("Mismatched brackets.")
+            open_brackets.pop()
+
+        pos += 1
+
+    if open_brackets:  # Append missing closing brackets
+        str_to_parse += "".join(reversed(open_brackets))
+
+    return str_to_parse

kronos/utils/fuzzy/_fuzzy_match.py

@@ -0,0 +1,149 @@
+from __future__ import annotations
+
+from enum import Enum
+from typing import TYPE_CHECKING, Any, Literal
+
+if TYPE_CHECKING:
+    from kronos.types import KeysLike
+
+from kronos.types._sentinel import Unset
+
+from ._string_similarity import SimilarityAlgo, string_similarity
+
+__all__ = ("fuzzy_match_keys",)
+
+HandleUnmatched = Literal["ignore", "raise", "remove", "fill", "force"]
+
+
+class HandleUnmatched(Enum):
+    """Strategy for handling unmatched keys in fuzzy_match_keys."""
+
+    IGNORE = "ignore"
+    """Keep original keys as-is."""
+
+    RAISE = "raise"
+    """Raise ValueError on unmatched keys."""
+
+    REMOVE = "remove"
+    """Remove unmatched original keys."""
+
+    FILL = "fill"
+    """Fill missing expected keys, keep unmatched original keys."""
+
+    FORCE = "force"
+    """Fill missing expected keys, remove unmatched original keys."""
+
+
+def fuzzy_match_keys(
+    d_: dict[str, Any],
+    keys: KeysLike,
+    /,
+    *,
+    similarity_algo: SimilarityAlgo = SimilarityAlgo.JARO_WINKLER,
+    similarity_threshold: float = 0.85,
+    fuzzy_match: bool = True,
+    handle_unmatched: HandleUnmatched = HandleUnmatched.IGNORE,
+    fill_value: Any = Unset,
+    fill_mapping: dict[str, Any] | None = None,
+    strict: bool = False,
+) -> dict[str, Any]:
+    """Validate and correct dict keys using fuzzy string matching.
+
+    Args:
+        d_: Input dictionary to validate.
+        keys: Expected keys (list, tuple, or object with .keys()).
+        similarity_algo: Algorithm for string similarity comparison.
+        similarity_threshold: Minimum similarity score (0.0-1.0).
+        fuzzy_match: Enable fuzzy matching for non-exact matches.
+        handle_unmatched: Strategy for unmatched keys (see HandleUnmatched).
+        fill_value: Default value for missing expected keys.
+        fill_mapping: Per-key values for specific missing keys.
+        strict: Raise if any expected keys remain unmatched.
+
+    Returns:
+        Dictionary with keys corrected to match expected keys.
+
+    Raises:
+        TypeError: If d_ is not dict or keys is None.
+        ValueError: If threshold out of range, or unmatched in RAISE mode.
+    """
+    # Input validation
+    if not isinstance(d_, dict):
+        raise TypeError("First argument must be a dictionary")
+    if keys is None:
+        raise TypeError("Keys argument cannot be None")
+    if not 0.0 <= similarity_threshold <= 1.0:
+        raise ValueError("similarity_threshold must be between 0.0 and 1.0")
+
+    # Extract expected keys
+    if isinstance(keys, (list, tuple)):
+        fields_set = set(keys)
+    elif hasattr(keys, "keys"):
+        fields_set = set(keys.keys())
+    else:
+        fields_set = set(keys)
+    if not fields_set:
+        return d_.copy()
+
+    corrected_out = {}
+    matched_expected = set()
+    matched_input = set()
+
+    # Pass 1: exact matches
+    for key in d_:
+        if key in fields_set:
+            corrected_out[key] = d_[key]
+            matched_expected.add(key)
+            matched_input.add(key)
+
+    # Pass 2: fuzzy matching
+    if fuzzy_match:
+        remaining_input = set(d_.keys()) - matched_input
+        remaining_expected = fields_set - matched_expected
+
+        for key in remaining_input:
+            if not remaining_expected:
+                break
+
+            matches = string_similarity(
+                key,
+                list(remaining_expected),
+                algorithm=similarity_algo,
+                threshold=similarity_threshold,
+                return_most_similar=True,
+            )
+
+            if matches:
+                match = matches if isinstance(matches, str) else matches[0] if matches else None
+                if match:
+                    corrected_out[match] = d_[key]
+                    matched_expected.add(match)
+                    matched_input.add(key)
+                    remaining_expected.remove(match)
+            elif handle_unmatched == HandleUnmatched.IGNORE:
+                corrected_out[key] = d_[key]
+
+    unmatched_input = set(d_.keys()) - matched_input
+    unmatched_expected = fields_set - matched_expected
+
+    if handle_unmatched == HandleUnmatched.RAISE and unmatched_input:
+        raise ValueError(f"Unmatched keys found: {unmatched_input}")
+
+    elif handle_unmatched == HandleUnmatched.IGNORE:
+        for key in unmatched_input:
+            corrected_out[key] = d_[key]
+
+    elif handle_unmatched in (HandleUnmatched.FILL, HandleUnmatched.FORCE):
+        for key in unmatched_expected:
+            corrected_out[key] = (
+                fill_mapping[key] if fill_mapping and key in fill_mapping else fill_value
+            )
+
+        if handle_unmatched == HandleUnmatched.FILL:
+            for key in unmatched_input:
+                corrected_out[key] = d_[key]
+
+    if strict and unmatched_expected:
+        raise ValueError(f"Missing required keys: {unmatched_expected}")
+
+    return corrected_out

kronos/utils/fuzzy/_string_similarity.py

@@ -0,0 +1,187 @@
+# Copyright (c) 2025 - 2026, HaiyangLi <quantocean.li at gmail dot com>
+# SPDX-License-Identifier: Apache-2.0
+
+"""String similarity algorithms with rapidfuzz backend.
+
+Provides multiple similarity metrics (Jaro-Winkler, Levenshtein, Hamming,
+cosine, SequenceMatcher) for fuzzy string matching and correction.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass
+from enum import Enum
+from typing import TYPE_CHECKING, Literal
+
+import rapidfuzz.distance as _rf_distance
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+__all__ = ("SimilarityAlgo", "string_similarity")
+
+
+class SimilarityAlgo(Enum):
+    """String similarity algorithms with associated scoring functions."""
+
+    JARO_WINKLER = "jaro_winkler"
+    LEVENSHTEIN = "levenshtein"
+    SEQUENCE_MATCHER = "sequence_matcher"
+    HAMMING = "hamming"
+    COSINE = "cosine"
+
+    @classmethod
+    def from_potential_valid(cls, v: SimilarityAlgo | str) -> SimilarityAlgo:
+        """Convert string to enum, passing through existing enum values."""
+        if isinstance(v, cls):
+            return v
+        try:
+            return cls(v)
+        except ValueError as e:
+            raise ValueError(f"Invalid similarity algorithm: {v}") from e
+
+    @staticmethod
+    def cosine_similarity(s1: str, s2: str) -> float:
+        """Character-set cosine similarity. Returns 0.0 if either string empty."""
+        if not s1 or not s2:
+            return 0.0
+        set1, set2 = set(s1), set(s2)
+        return len(set1 & set2) / ((len(set1) * len(set2)) ** 0.5)
+
+    @staticmethod
+    def hamming_similarity(s1: str, s2: str) -> float:
+        """Positional match ratio. Requires equal-length strings (else 0.0)."""
+        if not s1 or not s2 or len(s1) != len(s2):
+            return 0.0
+        return sum(c1 == c2 for c1, c2 in zip(s1, s2, strict=False)) / len(s1)
+
+    @staticmethod
+    def jaro_winkler_similarity(s: str, t: str, scaling: float = 0.1) -> float:
+        """Jaro-Winkler similarity with prefix boost (via rapidfuzz).
+
+        Args:
+            s: First string.
+            t: Second string.
+            scaling: Prefix weight, must be in [0, 0.25].
+
+        Returns:
+            Similarity score in [0, 1].
+
+        Raises:
+            ValueError: If scaling not in [0, 0.25].
+        """
+        if not 0 <= scaling <= 0.25:
+            raise ValueError("Scaling factor must be between 0 and 0.25")
+        return _rf_distance.JaroWinkler.similarity(s, t, prefix_weight=scaling)
+
+    @staticmethod
+    def sequence_matcher_similarity(s1: str, s2: str) -> float:
+        """Similarity via difflib.SequenceMatcher (longest contiguous match)."""
+        from difflib import SequenceMatcher
+
+        return SequenceMatcher(None, s1, s2).ratio()
+
+
+SIMILARITY_TYPE = Literal[
+    "jaro_winkler",
+    "levenshtein",
+    "sequence_matcher",
+    "hamming",
+    "cosine",
+]
+
+SimilarityFunc = Callable[[str, str], float]
+"""Type alias: (str, str) -> float similarity score."""
+
+
+@dataclass(frozen=True, slots=True)
+class MatchResult:
+    """String match with score and original index."""
+
+    word: str
+    score: float
+    index: int
+
+
+def string_similarity(
+    word: str,
+    correct_words: Sequence[str],
+    algorithm: SimilarityAlgo = SimilarityAlgo.JARO_WINKLER,
+    threshold: float = 0.0,
+    case_sensitive: bool = False,
+    return_most_similar: bool = False,
+) -> str | list[str] | None:
+    """Find similar strings using specified similarity algorithm.
+
+    Args:
+        word: The input string to find matches for
+        correct_words: List of strings to compare against
+        algorithm: Similarity algorithm to use
+        threshold: Minimum similarity score (0.0 to 1.0)
+        case_sensitive: Whether to consider case when matching
+        return_most_similar: Return only the most similar match
+
+    Returns:
+        Matching string(s) or None if no matches found
+
+    Raises:
+        ValueError: If correct_words is empty or threshold is invalid
+    """
+    if not correct_words:
+        raise ValueError("correct_words must not be empty")
+
+    if not 0.0 <= threshold <= 1.0:
+        raise ValueError("threshold must be between 0.0 and 1.0")
+
+    compare_word = str(word)
+    original_words = [str(w) for w in correct_words]
+
+    if not case_sensitive:
+        compare_word = compare_word.lower()
+        compare_words = [w.lower() for w in original_words]
+    else:
+        compare_words = original_words.copy()
+
+    algo_name: str | None = None
+    score_func: Callable[[str, str], float]
+    match algorithm:
+        case SimilarityAlgo.JARO_WINKLER:
+            algo_name = "jaro_winkler"
+            score_func = SimilarityAlgo.jaro_winkler_similarity
+        case SimilarityAlgo.LEVENSHTEIN:
+            algo_name = "levenshtein"
+            score_func = _rf_distance.Levenshtein.similarity
+        case SimilarityAlgo.SEQUENCE_MATCHER:
+            algo_name = "sequence_matcher"
+            score_func = SimilarityAlgo.sequence_matcher_similarity
+        case SimilarityAlgo.HAMMING:
+            algo_name = "hamming"
+            score_func = SimilarityAlgo.hamming_similarity
+        case SimilarityAlgo.COSINE:
+            algo_name = "cosine"
+            score_func = SimilarityAlgo.cosine_similarity
+        case _:
+            raise ValueError(f"Unsupported algorithm: {algorithm}")
+
+    results = []
+    for idx, (orig_word, comp_word) in enumerate(zip(original_words, compare_words, strict=False)):
+        if algo_name == "hamming" and len(comp_word) != len(compare_word):
+            continue  # Hamming requires equal length
+        score = score_func(compare_word, comp_word)  # type: ignore[operator]
+        if score >= threshold:
+            results.append(MatchResult(orig_word, score, idx))
+
+    if not results:
+        return None
+
+    results.sort(key=lambda x: (-x.score, x.index))  # Desc score, asc index
+
+    if case_sensitive:  # Keep only top-scoring matches
+        max_score = results[0].score
+        results = [r for r in results if r.score == max_score]
+
+    if return_most_similar:
+        return results[0].word
+
+    return [r.word for r in results]