typthon/Programs/_freeze_module.py

"""Python implementation of Programs/_freeze_module.c with strict type checking.

The pure Python implementation uses same functions and arguments as the C
implementation.

The generated byte code is slightly different because
compile() sets the PyCF_SOURCE_IS_UTF8 flag and objects have a
reference count > 1. Marshal adds the `FLAG_REF` flag and creates a
reference `hashtable`.

This version enforces strict type checking with runtime validation.
An exception is thrown if something is not typed correctly.
"""

import marshal
import sys
from typing import TextIO

header: str = "/* Auto-generated by Programs/_freeze_module.py */"


def _check_type(value, expected_type, param_name: str) -> None:
    """Enforce strict type checking at runtime.
    
    Args:
        value: The value to check
        expected_type: The expected type(s)
        param_name: Name of the parameter for error messages
        
    Raises:
        TypeError: If the value doesn't match the expected type
    """
    if not isinstance(value, expected_type):
        raise TypeError(
            f"Parameter '{param_name}' must be of type {expected_type.__name__}, "
            f"got {type(value).__name__}"
        )


def read_text(inpath: str) -> bytes:
    """Read text from a file and return as bytes.
    
    Args:
        inpath: Path to the input file
        
    Returns:
        The file contents as bytes
        
    Raises:
        TypeError: If inpath is not a string
    """
    _check_type(inpath, str, "inpath")
    
    with open(inpath, "rb") as f:
        result: bytes = f.read()
    
    return result


def compile_and_marshal(name: str, text: bytes) -> bytes:
    """Compile Python source code and marshal it.
    
    Args:
        name: Name identifier for the frozen module
        text: Python source code as bytes
        
    Returns:
        Marshalled bytecode
        
    Raises:
        TypeError: If parameters are not of correct types
    """
    _check_type(name, str, "name")
    _check_type(text, bytes, "text")
    
    filename: str = f"<frozen {name}>"
    # exec == Py_file_input
    code = compile(text, filename, "exec", optimize=0, dont_inherit=True)
    result: bytes = marshal.dumps(code)
    
    return result


def get_varname(name: str, prefix: str) -> str:
    """Generate a C variable name from a module name.
    
    Args:
        name: Module name
        prefix: Prefix for the variable name
        
    Returns:
        C-compatible variable name
        
    Raises:
        TypeError: If parameters are not strings
    """
    _check_type(name, str, "name")
    _check_type(prefix, str, "prefix")
    
    result: str = f"{prefix}{name.replace('.', '_')}"
    return result


def write_code(outfile: TextIO, marshalled: bytes, varname: str) -> None:
    """Write marshalled code as C array to output file.
    
    Args:
        outfile: Output file object
        marshalled: Marshalled bytecode
        varname: C variable name for the array
        
    Raises:
        TypeError: If parameters are not of correct types
    """
    _check_type(marshalled, bytes, "marshalled")
    _check_type(varname, str, "varname")
    
    # Check that outfile has a write method (duck typing for file-like objects)
    if not hasattr(outfile, 'write') or not callable(getattr(outfile, 'write')):
        raise TypeError(
            f"Parameter 'outfile' must be a file-like object with write() method"
        )
    
    data_size: int = len(marshalled)

    outfile.write(f"const unsigned char {varname}[] = {{\n")

    for n in range(0, data_size, 16):
        outfile.write("    ")
        outfile.write(",".join(str(i) for i in marshalled[n : n + 16]))
        outfile.write(",\n")
    outfile.write("};\n")


def write_frozen(outpath: str, inpath: str, name: str, marshalled: bytes) -> None:
    """Write frozen module to output file.
    
    Args:
        outpath: Path to output file
        inpath: Path to input file (unused but kept for API compatibility)
        name: Module name
        marshalled: Marshalled bytecode
        
    Raises:
        TypeError: If parameters are not of correct types
    """
    _check_type(outpath, str, "outpath")
    _check_type(inpath, str, "inpath")
    _check_type(name, str, "name")
    _check_type(marshalled, bytes, "marshalled")
    
    with open(outpath, "w") as outfile:
        outfile.write(header)
        outfile.write("\n")
        arrayname: str = get_varname(name, "_Py_M__")
        write_code(outfile, marshalled, arrayname)


def main() -> None:
    """Main entry point for the freeze_module program.
    
    Validates command-line arguments and processes the module freezing.
    
    Raises:
        SystemExit: If incorrect number of arguments provided
        TypeError: If any type validation fails
    """
    if len(sys.argv) != 4:
        sys.exit("need to specify the name, input and output paths\n")

    name: str = sys.argv[1]
    inpath: str = sys.argv[2]
    outpath: str = sys.argv[3]
    
    # Validate that arguments are strings
    _check_type(name, str, "name")
    _check_type(inpath, str, "inpath")
    _check_type(outpath, str, "outpath")

    text: bytes = read_text(inpath)
    marshalled: bytes = compile_and_marshal(name, text)
    write_frozen(outpath, inpath, name, marshalled)


if __name__ == "__main__":
    main()