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.
All function parameters must have type annotations.
"""

import inspect
import marshal
import sys
from typing import Any, Callable, TextIO, Type, TypeVar

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

F = TypeVar('F', bound=Callable[..., Any])


def strict_types(func: F) -> F:
    """Decorator that enforces type annotations on all parameters.
    
    Raises:
        TypeError: If any parameter lacks a type annotation
    """
    sig = inspect.signature(func)
    
    # Check that all parameters have annotations
    for param_name, param in sig.parameters.items():
        if param.annotation == inspect.Parameter.empty:
            raise TypeError(
                f"Function '{func.__name__}' parameter '{param_name}' "
                f"must have a type annotation"
            )
    
    # Check that return type is annotated
    if sig.return_annotation == inspect.Signature.empty:
        raise TypeError(
            f"Function '{func.__name__}' must have a return type annotation"
        )
    
    return func


def _check_type(value: Any, expected_type: 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):
        # Handle types that may not have __name__ attribute
        type_name = getattr(expected_type, '__name__', str(expected_type))
        raise TypeError(
            f"Parameter '{param_name}' must be of type {type_name}, "
            f"got {type(value).__name__}"
        )


@strict_types
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


@strict_types
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 == Ty_file_input
    code = compile(text, filename, "exec", optimize=0, dont_inherit=True)
    result: bytes = marshal.dumps(code)
    
    return result


@strict_types
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


@strict_types
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")


@strict_types
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, "_Ty_M__")
        write_code(outfile, marshalled, arrayname)


@strict_types
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]

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


if __name__ == "__main__":
    main()