SnowState - Mermaid Diagram Generator

snowstate_diagram.py

# ------------------------------------------------------------------------------------------------
# This script documents SnowState transitions by creating a mermaid diagram.
# You must be using SnowState's `add_transition` function to get output from this script.
#
# Arguments:
# ```
#   -i, --input_path: The input file or directory path to recursively search for transitions.
#   -o, --output_directory: The directory to output files to. Defaults to the current directory.
#   --html: Flag to output files as easily viewable HTML files instead of .mermaid files.
# ```
#
# Usage:
# ```
#   python snowstate_diagram.py -i <input_path> -o <optional_output_directory> --html
# ```
# ------------------------------------------------------------------------------------------------

import argparse
import glob
import json
import os
import re
from typing import TypedDict

class ParsedTransition(TypedDict):
  transition_name: str
  transition_from: str | list[str]
  transition_to: str

FILE_TYPE = "gml"

# -----------------------------------------------------------------------------------
# Arguments
# -----------------------------------------------------------------------------------
parser = argparse.ArgumentParser(description="Find all .gd files in a directory.")
parser.add_argument("-i", "--input_path", type=str, help="The input file or directory path to recursively search for transitions.", required=True)
parser.add_argument("-o", "--output_directory", type=str, help="The directory to output files to. Defaults to the current directory.", default=".")
parser.add_argument("--html", action='store_true', help="Flag to output files as easily viewable HTML files instead of .mermaid files.")

# -----------------------------------------------------------------------------------
# Utility Functions
# -----------------------------------------------------------------------------------
def find_fsm_transition_funcs_in_text(text: str) -> list[str]:
  matches = re.finditer(r'\.add_transition', text)

  results = []
  for match in matches:
    start_index = match.end() + 1 # +2 to skip the first parenthesis
    current_index = start_index
    end_index = None
    parenthesis_nested_depth = 1 # Starting at 1 because we're already inside a parenthesis

    while end_index is None:
      character = text[current_index]
      if character == "(": parenthesis_nested_depth += 1
      if character == ")": parenthesis_nested_depth -= 1

      if parenthesis_nested_depth == 0:
        end_index = current_index
        break

      current_index += 1
    results.append(text[start_index:end_index])
  return results

def parse_transition_arguments(arguments: str) -> ParsedTransition | None:
  args = ["", "", ""]
  arg_index = 0
  ignore_commas = False

  # Manually loop through each character to handle commas inside of arrays
  for character in arguments:
    if character == "[": ignore_commas = True
    if character == "]": ignore_commas = False

    if character == "," and not ignore_commas:
      if arg_index == 2: break

      arg_index += 1
      continue

    args[arg_index] += character
  
  # Check if transition from is an array
  transition_from = None
  transition_from_is_string_array = "[" in args[1] and "]" in args[1]

  if transition_from_is_string_array:
    transition_from = json.loads(args[1])
  else:
    transition_from = args[1]
    transition_from = transition_from.replace("\"", "").strip()

  return {
    "transition_name": args[0].replace("\"", "").strip(),
    "transition_from": transition_from,
    "transition_to": args[2].replace("\"", "").strip(),
  }

def find_files_by_filetype(directory: str, filetype: str):
  path = os.path.join(directory, "**", f"*.{filetype}")
  return glob.glob(path, recursive=True)

def read_file(file_path: str):
  with open(file_path, "r") as file:
    return file.read()

def validate_transition_args(args: dict) -> bool:
  required_keys = ["transition_name", "transition_from", "transition_to"]
  return all(key in args for key in required_keys)

def create_html_output(diagram: str):
  return f"""
<!DOCTYPE html>
<div class="mermaid">
{diagram}
</div>
<script src="https://cdn.jsdelivr.net/npm/mermaid/dist/mermaid.min.js"></script>
<script>
  mermaid.initialize({{ startOnLoad: true }});
</script>
"""

# -----------------------------------------------------------------------------------
# Main
# -----------------------------------------------------------------------------------
if __name__ == "__main__":
  args = parser.parse_args()
  transitions_by_file: dict[str, list[str]] = {}

  # Validate input path
  if not os.path.exists(args.input_path):
    print(f"Input path {args.input_path} does not exist.")
    exit(1)

  # Create output directory if it doesn't exist
  if not os.path.exists(args.output_directory):
    os.makedirs(args.output_directory)

  # Account for input path being a directory or file
  if os.path.isdir(args.input_path):
    file_paths = find_files_by_filetype(args.input_path, FILE_TYPE)
    print(f"Discovered {len(file_paths)} {FILE_TYPE} files...")
  else:
    file_paths = [args.input_path]

  # Find transitions in each file
  for file_path in file_paths:
    folder_name = os.path.basename(os.path.dirname(file_path))
    file_text = read_file(file_path)

    discovered_transitions = find_fsm_transition_funcs_in_text(file_text)
    if len(discovered_transitions) == 0: continue
    transitions_by_file[folder_name] = discovered_transitions

  # Parse found transitions
  parsed_transitions_by_file: dict[str, list[ParsedTransition]] = {}

  for folder_name, file_transitions in transitions_by_file.items():
    # Loop through each transition in the file
    for transition in file_transitions:
      parsed_transition = parse_transition_arguments(transition)
      if parsed_transition is None: continue

      if folder_name not in parsed_transitions_by_file.keys(): parsed_transitions_by_file[folder_name] = []
      parsed_transitions_by_file[folder_name].append(parsed_transition)

  transition_files_count = len(parsed_transitions_by_file)
  print(f"Found {transition_files_count} files with transitions...")

  # Output to mermaid format
  for folder_name, parsed_transitions in parsed_transitions_by_file.items():
    output = f"---\ntitle: {folder_name}\n---\nstateDiagram-v2\n"

    for parsed_transition in parsed_transitions:
      transition_name = parsed_transition["transition_name"]
      transition_from = parsed_transition["transition_from"]
      transition_to = parsed_transition["transition_to"]

      if isinstance(transition_from, list):
        for from_state in transition_from:
          output += f"\t{from_state} --> {transition_to} : {transition_name}\n"
      else:
        output += f"\t{transition_from} --> {transition_to} : {transition_name}\n"

    if args.html:
      html_output = create_html_output(output)
      with open(os.path.join(args.output_directory, f"{folder_name}.html"), "w") as file:
        file.write(html_output)
    else:
      with open(os.path.join(args.output_directory, f"{folder_name}.mermaid"), "w") as file:
        file.write(output)

  print(f"Output {transition_files_count} files to { os.path.abspath(args.output_directory) }")