exo/format.py

#!/usr/bin/env python
import re
import subprocess
import sys
import os
import fnmatch

DEBUG_PATTERN = re.compile(r'^(\s*)(if\s+DEBUG\s*>=?\s*\d+\s*:.+)$', re.MULTILINE)
PLACEHOLDER = "###DEBUG_PLACEHOLDER###"

# Add ignore patterns here
IGNORE_PATTERNS = [
  '.venv/*',
  'setup.py',
  '*helpers.py',
  '*node_service_pb2.py',
  '*node_service_pb2_grpc.py',
]


def should_ignore(file_path):
  for pattern in IGNORE_PATTERNS:
    if fnmatch.fnmatch(file_path, pattern):
      return True
  return False


def preserve_debug_lines(content):
  def replace(match):
    indent, line = match.groups()
    return f"{indent}{PLACEHOLDER}{line.strip()}"

  return DEBUG_PATTERN.sub(replace, content)


def restore_debug_lines(content):
  return re.sub(f"^(\\s*){PLACEHOLDER}(.+)$", r"\1\2", content, flags=re.MULTILINE)


def adjust_indentation(content):
  lines = content.split('\n')
  adjusted_lines = []
  for line in lines:
    if line.strip() and not line.startswith(PLACEHOLDER):
      indent = len(line) - len(line.lstrip())
      new_indent = ' ' * (indent // 2)
      adjusted_lines.append(new_indent + line.lstrip())
    else:
      adjusted_lines.append(line)
  return '\n'.join(adjusted_lines)


def process_file(file_path, process_func):
  with open(file_path, 'r') as file:
    content = file.read()

  modified_content = process_func(content)

  if content != modified_content:
    with open(file_path, 'w') as file:
      file.write(modified_content)


def run_black(target):
  # Convert ignore patterns to Black's --extend-exclude format
  exclude_patterns = '|'.join(f'({pattern.replace("*", ".*")})' for pattern in IGNORE_PATTERNS)
  command = ["black", "--line-length", "200", "--extend-exclude", exclude_patterns, target]
  subprocess.run(command, check=True)


def format_files(target):
  if os.path.isfile(target):
    files = [target] if not should_ignore(target) else []
  elif os.path.isdir(target):
    files = []
    for root, _, filenames in os.walk(target):
      for filename in filenames:
        if filename.endswith('.py'):
          file_path = os.path.join(root, filename)
          if not should_ignore(file_path):
            files.append(file_path)
  else:
    print(f"Error: {target} is not a valid file or directory")
    return

  # Preserve debug lines
  for file in files:
    process_file(file, preserve_debug_lines)

  # Run Black
  run_black(target)

  # Adjust indentation and restore debug lines
  for file in files:
    process_file(file, adjust_indentation)
    process_file(file, restore_debug_lines)


def main():
  if len(sys.argv) < 2:
    print("Usage: python format.py <directory_or_file>")
    sys.exit(1)

  target = sys.argv[1]
  format_files(target)
  print("Formatting completed.")


if __name__ == "__main__":
  main()