Source code for astropy.utils.codegen

# Licensed under a 3-clause BSD style license - see LICENSE.rst
"""Utilities for generating new Python code at runtime."""

import inspect
import keyword
import os
import re
import textwrap

from .introspection import find_current_module

__all__ = ["make_function_with_signature"]


_ARGNAME_RE = re.compile(r"^[A-Za-z][A-Za-z_]*")
"""
Regular expression used my make_func which limits the allowed argument
names for the created function.  Only valid Python variable names in
the ASCII range and not beginning with '_' are allowed, currently.
"""




def make_function_with_signature(
    func, args=(), kwargs={}, varargs=None, varkwargs=None, name=None
):
    """
    Make a new function from an existing function but with the desired
    signature.

    The desired signature must of course be compatible with the arguments
    actually accepted by the input function.

    The ``args`` are strings that should be the names of the positional
    arguments.  ``kwargs`` can map names of keyword arguments to their
    default values.  It may be either a ``dict`` or a list of ``(keyword,
    default)`` tuples.

    If ``varargs`` is a string it is added to the positional arguments as
    ``*<varargs>``.  Likewise ``varkwargs`` can be the name for a variable
    keyword argument placeholder like ``**<varkwargs>``.

    If not specified the name of the new function is taken from the original
    function.  Otherwise, the ``name`` argument can be used to specify a new
    name.

    Note, the names may only be valid Python variable names.
    """
    pos_args = []
    key_args = []

    if isinstance(kwargs, dict):
        iter_kwargs = kwargs.items()
    else:
        iter_kwargs = iter(kwargs)

    # Check that all the argument names are valid
    for item in (*args, *iter_kwargs):
        if isinstance(item, tuple):
            argname = item[0]
            key_args.append(item)
        else:
            argname = item
            pos_args.append(item)

        if keyword.iskeyword(argname) or not _ARGNAME_RE.match(argname):
            raise SyntaxError(f"invalid argument name: {argname}")

    for item in (varargs, varkwargs):
        if item is not None:
            if keyword.iskeyword(item) or not _ARGNAME_RE.match(item):
                raise SyntaxError(f"invalid argument name: {item}")

    def_signature = [", ".join(pos_args)]

    if varargs:
        def_signature.append(f", *{varargs}")

    call_signature = def_signature[:]

    if name is None:
        name = func.__name__

    global_vars = {f"__{name}__func": func}
    local_vars = {}
    # Make local variables to handle setting the default args
    for idx, item in enumerate(key_args):
        key, value = item
        default_var = f"_kwargs{idx}"
        local_vars[default_var] = value
        def_signature.append(f", {key}={default_var}")
        call_signature.append(f", {key}={key}")

    if varkwargs:
        def_signature.append(f", **{varkwargs}")
        call_signature.append(f", **{varkwargs}")

    def_signature = "".join(def_signature).lstrip(", ")
    call_signature = "".join(call_signature).lstrip(", ")

    mod = find_current_module(2)
    frm = inspect.currentframe().f_back

    if mod:
        filename = mod.__file__
        modname = mod.__name__
        if filename.endswith(".pyc"):
            filename = os.path.splitext(filename)[0] + ".py"
    else:
        filename = "<string>"
        modname = "__main__"

    num_blank_lines = func.__code__.co_firstlineno - 1
    blank_lines = "\n" * num_blank_lines

    # The lstrip is in case there were *no* positional arguments (a rare case)
    # in any context this will actually be used...
    template = textwrap.dedent(
        f"""{blank_lines}\
    def {name}({def_signature}):
        return __{name}__func({call_signature})
    """
    )

    code = compile(template, filename, "single")

    eval(code, global_vars, local_vars)

    new_func = local_vars[name]
    new_func.__module__ = modname
    new_func.__doc__ = func.__doc__

    return new_func