Feature: Added move query functionality to Heihachi frame data bot

README.md

@@ -57,6 +57,7 @@ The bot supports the following slash commands -
 | Command | Description |
 | --- | --- |
 | `/fd <character> <move>` | Get frame data of a particular character's move |
+| `/ms <character> <condition> <frames> <situation>` | Find a character's moves that match a particular frame scenario |
 | `/<character>` | Get frame data for a particular character's move |
 | `/feedback <message>` | Send feedback to the bot owner |
 | `/help` | Get help on the bot's usage |

src/framedb/const.py

@@ -1,5 +1,5 @@
 import enum
-from typing import Dict, List
+from typing import Callable, Dict, List
 
 NUM_CHARACTERS = 34
 
@@ -95,6 +95,12 @@ class MoveType(enum.Enum):
     HS = "Heat Smash"
 
 
+class FrameSituation(enum.Enum):
+    STARTUP = "startup"
+    BLOCK = "block"
+    HIT = "hit"
+
+
 MOVE_TYPE_ALIAS: Dict[MoveType, List[str]] = {
     MoveType.RA: ["ra", "rage_art", "rageart", "rage art"],
     MoveType.T: ["screw", "t!", "t", "screws", "tor", "tornado"],
@@ -154,3 +160,11 @@ class MoveType(enum.Enum):
     "4\ufe0f\u20e3",
     "5\ufe0f\u20e3",
 ]
+
+CONDITION_MAP: Dict[str, Callable[[int, int], bool]] = {
+    ">": lambda x, y: x > y,
+    ">=": lambda x, y: x >= y,
+    "<": lambda x, y: x < y,
+    "<=": lambda x, y: x <= y,
+    "==": lambda x, y: x == y,
+}

src/framedb/framedb.py

@@ -1,5 +1,6 @@
 import logging
 import os
+import re
 from typing import Dict, List
 
 import requests
@@ -8,7 +9,7 @@
 from fast_autocomplete import AutoComplete
 
 from .character import Character, Move
-from .const import CHARACTER_ALIAS, MOVE_TYPE_ALIAS, REPLACE, CharacterName, MoveType
+from .const import CHARACTER_ALIAS, CONDITION_MAP, MOVE_TYPE_ALIAS, REPLACE, CharacterName, FrameSituation, MoveType
 from .frame_service import FrameService
 
 MATCH_SCORE_CUTOFF = 95
@@ -103,6 +104,21 @@ def _is_command_in_alt(move_query: str, move: Move) -> bool:
                 return True
         return False
 
+    @staticmethod
+    def _sanitize_frame_data(frame_data: str) -> int | None:
+        """Removes bells and whistles from a move's frame data result'"""
+
+        # Step 1: Remove any leading non-numeric characters (like 'i' or ',')
+        frame_data = re.sub(r"^[^-\d]+", "", frame_data)
+
+        # Step 2: Match the first number or range (e.g., "-5", "5", "5~10")
+        match = re.match(r"([-+]?\d+)", frame_data)
+
+        # Step 3: If a match is found, return the first number (and remove any '+' sign)
+        if match:
+            return int(match.group(1))
+        return None
+
     def get_move_by_input(self, character: CharacterName, input_query: str) -> Move | None:
         """Given an input move query for a known character, retrieve the move from the database."""
 
@@ -130,6 +146,46 @@ def get_move_by_input(self, character: CharacterName, input_query: str) -> Move
         # couldn't match anything :-(
         return None
 
+    def get_move_by_frame(
+        self, character: CharacterName, condition: str, frame_value: int, situation: FrameSituation
+    ) -> List[Move]:
+        """Given a frame value query for a known character, retrieve the moves from the database that matches that frame value."""
+
+        character_movelist = self.frames[character].movelist.values()
+        result = []
+
+        # Get the comparison function based on the condition
+        compare_func = CONDITION_MAP.get(condition.strip())
+        if compare_func is None:
+            raise ValueError(f"Unsupported condition: {condition}")
+
+        match situation:
+            case FrameSituation.STARTUP:
+                result = [
+                    entry
+                    for entry in character_movelist
+                    if entry.startup.strip() != ""  # Ignore moves with no frame data
+                    if (sanitized_value := FrameDb._sanitize_frame_data(entry.startup)) is not None
+                    if compare_func(sanitized_value, frame_value)
+                ]
+            case FrameSituation.BLOCK:
+                result = [
+                    entry
+                    for entry in character_movelist
+                    if entry.on_block.strip() != ""  # Ignores moves with no frame data
+                    if (sanitized_value := FrameDb._sanitize_frame_data(entry.on_block)) is not None
+                    if compare_func(sanitized_value, frame_value)
+                ]
+            case FrameSituation.HIT:
+                result = [
+                    entry
+                    for entry in character_movelist
+                    if entry.on_hit.strip() != ""  # Ignores moves with no frame data
+                    if (sanitized_value := FrameDb._sanitize_frame_data(entry.on_hit)) is not None
+                    if compare_func(sanitized_value, frame_value)
+                ]
+        return result
+
     def get_moves_by_move_name(self, character: CharacterName, move_name_query: str) -> List[Move]:
         """
         Gets a list of moves that match a move name query, by comparing the move name and its aliases

src/framedb/tests/conftest.py

@@ -0,0 +1,7 @@
+import os
+import sys
+
+import pytest
+
+# Add the project root directory to sys.path
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../..")))

src/framedb/tests/test_framedb.py

@@ -2,7 +2,7 @@
 
 from frame_service.json_directory.tests.test_json_directory import json_directory
 from framedb import FrameDb, FrameService
-from framedb.const import CharacterName, MoveType
+from framedb.const import CharacterName, FrameSituation, MoveType
 
 
 @pytest.fixture
@@ -85,3 +85,223 @@ def test_search_move() -> None:
 def test_all_autocomplete_words_match() -> None:
     "Test that all words in the autocomplete list can be matched to a CharacterName"
     pass
+
+
+def test_sanitize_frame_data(frameDb: FrameDb) -> None:
+    assert frameDb._sanitize_frame_data("i59~61") == 59
+    assert frameDb._sanitize_frame_data(",i13,14,15") == 13
+    assert frameDb._sanitize_frame_data("i13~14,i25 i35 i39 i42") == 13
+    assert frameDb._sanitize_frame_data("-5") == -5
+    assert frameDb._sanitize_frame_data("+5") == 5
+    assert frameDb._sanitize_frame_data("+5~10") == 5
+    assert frameDb._sanitize_frame_data("i16") == 16
+    assert frameDb._sanitize_frame_data("i5~8") == 5
+    assert frameDb._sanitize_frame_data("-5~-10") == -5
+    assert frameDb._sanitize_frame_data("+67a (+51)") == 67
+    assert frameDb._sanitize_frame_data("invalid data") is None
+    assert frameDb._sanitize_frame_data("") is None
+
+
+def test_get_move_by_frame_on_block(frameDb: FrameDb) -> None:
+    # == Case
+    returned_moves = frameDb.get_move_by_frame(CharacterName.RAVEN, "==", -5, FrameSituation.BLOCK)
+    move_ids = list(map(lambda move: move.id, returned_moves))
+    expected_move_ids = [
+        "Raven-BT.4",
+        "Raven-BT.f+2",
+        "Raven-BT.f+3",
+        "Raven-d+1",
+        "Raven-FC.1",
+        "Raven-H.BT.4,F",
+    ]
+    assert set(move_ids) == set(expected_move_ids)
+
+    # > Case
+    returned_moves = frameDb.get_move_by_frame(CharacterName.RAVEN, ">", 5, FrameSituation.BLOCK)
+    move_ids = list(map(lambda move: move.id, returned_moves))
+    expected_move_ids = ["Raven-(Back to wall).b,b,UB", "Raven-b+1", "Raven-H.f,f,F+3,4", "Raven-H.2+3"]
+    assert set(move_ids) == set(expected_move_ids)
+
+    # >= Case
+    returned_moves = frameDb.get_move_by_frame(CharacterName.RAVEN, ">=", 5, FrameSituation.BLOCK)
+    move_ids = list(map(lambda move: move.id, returned_moves))
+    expected_move_ids = [
+        "Raven-H.1+2,F",
+        "Raven-H.f+1+2,F",
+        "Raven-H.SZN.2,F",
+        "Raven-H.ws3+4,F",
+        "Raven-(Back to wall).b,b,UB",
+        "Raven-b+1",
+        "Raven-H.f,f,F+3,4",
+        "Raven-H.2+3",
+    ]
+    assert set(move_ids) == set(expected_move_ids)
+
+    # < Case
+    returned_moves = frameDb.get_move_by_frame(CharacterName.RAVEN, "<", -25, FrameSituation.BLOCK)
+    move_ids = list(map(lambda move: move.id, returned_moves))
+    expected_move_ids = ["Raven-3~4,F", "Raven-b+2,2,1+2", "Raven-BT.3,4,4,F", "Raven-BT.d+3", "Raven-BT.f+3+4,F"]
+    assert set(move_ids) == set(expected_move_ids)
+
+    # <= Case
+    returned_moves = frameDb.get_move_by_frame(CharacterName.RAVEN, "<=", -25, FrameSituation.BLOCK)
+    move_ids = list(map(lambda move: move.id, returned_moves))
+    expected_move_ids = [
+        "Raven-3~4,F",
+        "Raven-b+2,2,1+2",
+        "Raven-BT.3,4,4,F",
+        "Raven-BT.d+3",
+        "Raven-BT.f+3+4,F",
+        "Raven-uf+3+4,3+4",
+        "Raven-b+4,B+4~3,3+4",
+    ]
+    assert set(move_ids) == set(expected_move_ids)
+
+
+def test_get_move_by_frame_on_hit(frameDb: FrameDb) -> None:
+    # == Case
+    returned_moves = frameDb.get_move_by_frame(CharacterName.RAVEN, "==", 12, FrameSituation.HIT)
+    move_ids = list(map(lambda move: move.id, returned_moves))
+    expected_move_ids = [
+        "Raven-b+2,2,3",
+        "Raven-df+3",
+        "Raven-SZN.1~F",
+    ]
+    assert set(move_ids) == set(expected_move_ids)
+
+    # > Case
+    returned_moves = frameDb.get_move_by_frame(CharacterName.RAVEN, ">", 45, FrameSituation.HIT)
+    move_ids = list(map(lambda move: move.id, returned_moves))
+    expected_move_ids = [
+        "Raven-H.ws3+4,F",
+        "Raven-ws3,2",
+        "Raven-H.f,f,F+3,4",
+        "Raven-H.ws3,2",
+    ]
+    assert set(move_ids) == set(expected_move_ids)
+
+    # >= Case
+    returned_moves = frameDb.get_move_by_frame(CharacterName.RAVEN, ">=", 44, FrameSituation.HIT)
+    move_ids = list(map(lambda move: move.id, returned_moves))
+    expected_move_ids = [
+        "Raven-H.1+2,F",
+        "Raven-BT.4",
+        "Raven-H.ws3+4,F",
+        "Raven-ws3,2",
+        "Raven-H.f,f,F+3,4",
+        "Raven-H.ws3,2",
+    ]
+    assert set(move_ids) == set(expected_move_ids)
+
+    # < Case
+    returned_moves = frameDb.get_move_by_frame(CharacterName.RAVEN, "<", -5, FrameSituation.HIT)
+    move_ids = list(map(lambda move: move.id, returned_moves))
+    expected_move_ids = ["Raven-FC.3", "Raven-3~4,F", "Raven-UB,b,3+4", "Raven-BT.f+3+4,F", "Raven-b+1+3,P"]
+    assert set(move_ids) == set(expected_move_ids)
+
+    # <= Case
+    returned_moves = frameDb.get_move_by_frame(CharacterName.RAVEN, "<=", -5, FrameSituation.HIT)
+    move_ids = list(map(lambda move: move.id, returned_moves))
+    expected_move_ids = [
+        "Raven-FC.3",
+        "Raven-3~4,F",
+        "Raven-UB,b,3+4",
+        "Raven-BT.f+3+4,F",
+        "Raven-b+1+3,P",
+        "Raven-FC.df+3+4",
+        "Raven-H.FC.df+3+4",
+    ]
+    assert set(move_ids) == set(expected_move_ids)
+
+
+def test_get_move_by_frame_startup(frameDb: FrameDb) -> None:
+    # == Case
+    returned_moves = frameDb.get_move_by_frame(CharacterName.RAVEN, "==", 11, FrameSituation.STARTUP)
+    move_ids = list(map(lambda move: move.id, returned_moves))
+    expected_move_ids = [
+        "Raven-d+2",
+        "Raven-df+1+4",
+        "Raven-FC.2",
+        "Raven-ws4",
+    ]
+    assert set(move_ids) == set(expected_move_ids)
+
+    # > Case
+    returned_moves = frameDb.get_move_by_frame(CharacterName.RAVEN, ">", 26, FrameSituation.STARTUP)
+    move_ids = list(map(lambda move: move.id, returned_moves))
+    expected_move_ids = [
+        "Raven-b+3",
+        "Raven-u+3",
+        "Raven-u+3,3",
+        "Raven-u+3,3,3",
+        "Raven-b+1+2",
+        "Raven-db+3",
+        "Raven-uf+3",
+        "Raven-(Back to wall).b,b,UB",
+    ]
+    assert set(move_ids) == set(expected_move_ids)
+
+    # >= Case
+    returned_moves = frameDb.get_move_by_frame(CharacterName.RAVEN, ">=", 26, FrameSituation.STARTUP)
+    move_ids = list(map(lambda move: move.id, returned_moves))
+    expected_move_ids = [
+        "Raven-H.f,f,F+3,4",
+        "Raven-f,f,F+3",
+        "Raven-b+3",
+        "Raven-u+3",
+        "Raven-u+3,3",
+        "Raven-u+3,3,3",
+        "Raven-b+1+2",
+        "Raven-db+3",
+        "Raven-uf+3",
+        "Raven-(Back to wall).b,b,UB",
+    ]
+    assert set(move_ids) == set(expected_move_ids)
+
+    # < Case
+    returned_moves = frameDb.get_move_by_frame(CharacterName.RAVEN, "<", 10, FrameSituation.STARTUP)
+    move_ids = list(map(lambda move: move.id, returned_moves))
+    expected_move_ids = [
+        "Raven-BT.1",
+        "Raven-BT.1,4",
+    ]
+    assert set(move_ids) == set(expected_move_ids)
+
+    # <= Case
+    returned_moves = frameDb.get_move_by_frame(CharacterName.RAVEN, "<=", 10, FrameSituation.STARTUP)
+    move_ids = list(map(lambda move: move.id, returned_moves))
+    expected_move_ids = [
+        "Raven-1",
+        "Raven-1,2",
+        "Raven-1,2,3+4",
+        "Raven-1,2,4",
+        "Raven-2",
+        "Raven-2,3",
+        "Raven-2,4",
+        "Raven-BT.1",
+        "Raven-BT.1,4",
+        "Raven-BT.2",
+        "Raven-BT.2,1",
+        "Raven-BT.2,1~F",
+        "Raven-BT.2,2",
+        "Raven-BT.2,2,1",
+        "Raven-BT.2,2,3+4",
+        "Raven-BT.3",
+        "Raven-BT.3,4",
+        "Raven-BT.3,4,3",
+        "Raven-BT.3,4,4",
+        "Raven-BT.3,4,4,F",
+        "Raven-BT.d+1",
+        "Raven-d+1",
+        "Raven-FC.1",
+        "Raven-H.BT.2,2,1",
+    ]
+
+    assert set(move_ids) == set(expected_move_ids)
+
+
+def test_get_move_by_frame_no_results(frameDb: FrameDb) -> None:
+    returned_moves = frameDb.get_move_by_frame(CharacterName.RAVEN, "<", 1, FrameSituation.STARTUP)
+    move_ids = list(map(lambda move: move.id, returned_moves))
+    expected_move_ids: list[str] = []
+    assert set(move_ids) == set(expected_move_ids)