<filename>astropy/utils/codegen.py
# -*- coding: utf-8 -*-
# Licensed under a 3-clause BSD style license - see LICENSE.rst
"""Utilities for generating new Python code at runtime."""


import inspect
import itertools
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 itertools.chain(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(', {0}={0}'.format(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__'

    # Subtract 2 from the line number since the length of the template itself
    # is two lines.  Therefore we have to subtract those off in order for the
    # pointer in tracebacks from __{name}__func to point to the right spot.
    lineno = frm.f_lineno - 2

    # The lstrip is in case there were *no* positional arguments (a rare case)
    # in any context this will actually be used...
    template = textwrap.dedent("""{0}\
    def {name}({sig1}):
        return __{name}__func({sig2})
    """.format('\n' * lineno, name=name, sig1=def_signature,
               sig2=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
