Python script to re/calculate EIP-7201 namespaced storage locations in your project. Run with `uv run eip7201.py <dir>`

eip7201.py

#!/usr/bin/env python3
# /// script
# requires-python = ">=3.7"
# dependencies = [
#     "pycryptodome>=3.17.0",
# ]
# ///

"""
This script processes Solidity (.sol) files to replace lines following
`///@custom:storage-location eip7201:<Namespace>` with a computed keccak256 slot.
It correctly removes multi-line definitions when they span multiple lines.

See https://eips.ethereum.org/EIPS/eip-7201#customstorage-location for information on EIP-7201.

Usage:
    python eip7201.py <directory>
"""

import sys
import os
import re
from Crypto.Hash import keccak


def keccak256(data: bytes) -> bytes:
    """
    Return the Keccak-256 digest of the given bytes using pycryptodome.
    """
    k = keccak.new(digest_bits=256)
    k.update(data)
    return k.digest()


def to_screaming_snake_case(identifier: str) -> str:
    """
    Convert 'MyCompany.MyCoolStorage' -> 'MY_COMPANY_MY_COOL_STORAGE'.
    Replaces periods with underscores, inserts underscores before uppercase letters,
    and uppercases everything.
    """
    # Replace periods with underscores
    identifier = identifier.replace('.', '_')
    # Insert underscores before uppercase letters
    identifier = re.sub(r'(?<!^)(?=[A-Z])', '_', identifier)
    # Convert to uppercase
    return identifier.upper()


def compute_slot_hex(namespace_id: str) -> str:
    """
    Implements:
      keccak256(abi.encode(uint256(keccak256(bytes(id))) - 1)) & ~bytes32(uint256(0xff))
    and returns a 0x-prefixed, 64-hex-digit string (the masked 32-byte value).
    """
    # 1) keccak256(bytes(namespace_id))
    ns_hash = keccak256(namespace_id.encode('utf-8'))
    ns_int = int.from_bytes(ns_hash, byteorder='big')

    # 2) subtract 1 (mod 2^256)
    ns_int_minus_one = (ns_int - 1) % (1 << 256)

    # 3) re-encode
    ns_minus_one_32 = ns_int_minus_one.to_bytes(32, byteorder='big')

    # 4) keccak256(...) again
    slot_hash = keccak256(ns_minus_one_32)
    slot_int = int.from_bytes(slot_hash, byteorder='big')

    # 5) zero out the last byte ( ~0xFF )
    slot_int_masked = slot_int & ~0xFF

    # 6) convert to hex string
    final_32_bytes = slot_int_masked.to_bytes(32, byteorder='big')
    return "0x" + final_32_bytes.hex()


def replace_storage_location_line(lines):
    """
    For each line matching:
       ///@custom:storage-location eip7201:SomeNamespace
    parse out SomeNamespace, compute the line to replace the *next* line with,
    and do the replacement.

    Handles multi-line definitions by checking if the next line ends with `=`.
    Also removes placeholder lines that follow the constant definition.
    """
    i = 0
    out_lines = []
    while i < len(lines):
        line = lines[i]

        # Regex for: ///@custom:storage-location eip7201:IDENTIFIER
        match = re.match(r'^\s*///@custom:storage-location\s+eip7201:([A-Za-z0-9_.]+)', line)
        if match:
            namespace_id = match.group(1)  # e.g. 'MyCompany.MyCoolStorage'
            out_lines.append(line)  # Keep the comment line

            # Check if the next line is a `bytes32 constant` declaration
            if i + 1 < len(lines):
                next_line = lines[i + 1]
                if re.search(r'=\s*$', next_line):  # If it ends with '=', it's multi-line
                    i += 3  # Skip the next line, the value line, and the duplicate line
                else:
                    i += 2  # Skip both the next line and the placeholder line

                # Grab indentation from the first replaced line
                indent_match = re.match(r'^(\s*)', next_line)
                indentation = indent_match.group(1) if indent_match else ''

                # Convert namespace to SCREAMING_SNAKE_CASE
                snake = to_screaming_snake_case(namespace_id)
                # Compute the slot
                slot_hex = compute_slot_hex(namespace_id)
                # Replacement line
                replacement = f"{indentation}bytes32 constant {snake}_SLOT = {slot_hex};\n"

                # Append the new slot definition
                out_lines.append(replacement)
                continue  # Skip processing the removed lines

        out_lines.append(line)
        i += 1

    return out_lines


def process_file(file_path: str) -> None:
    """
    Reads the Solidity (.sol) file, applies replacements, and writes back the result.
    """
    if not file_path.endswith('.sol'):
        return
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            lines = f.readlines()
    except (FileNotFoundError, OSError) as e:
        print(f"Error reading file '{file_path}': {e}")
        return

    new_lines = replace_storage_location_line(lines)

    try:
        with open(file_path, 'w', encoding='utf-8') as f:
            f.writelines(new_lines)
    except OSError as e:
        print(f"Error writing file '{file_path}': {e}")


def process_directory(directory: str) -> None:
    """
    Recursively processes all Solidity (.sol) files in the given directory.
    """
    for root, _, files in os.walk(directory):
        for file in files:
            if file.endswith('.sol'):
                file_path = os.path.join(root, file)
                print(f"Processing {file_path}...")
                process_file(file_path)


def main():
    if len(sys.argv) != 2:
        print("Usage: eip7201.py [directory]")
        sys.exit(1)

    directory = sys.argv[1]
    if not os.path.isdir(directory):
        print(f"Error: '{directory}' is not a valid directory.")
        sys.exit(1)

    process_directory(directory)
    print(f"Finished processing .sol files in '{directory}'.")


if __name__ == "__main__":
    main()