<commit_before>'''
A library of expressions that can be composed of existing expressions.
'''

from .udf import Function
from raco.expression import *


def lookup(function_name, num_args):
    func = EXPRESSIONS.get(function_name.lower())
    if hasattr(func, '__call__'):
        return func(num_args)
    if isinstance(func, dict):
        return func.get(num_args)
    return func


def create_nested_binary(num_args, func):
    if num_args < 2:
        return None

    def make_binary(num_args):
        if num_args == 2:
            return func(
                NamedAttributeRef('x%d' % num_args),
                NamedAttributeRef('x%d' % (num_args - 1)))
        return func(
            NamedAttributeRef('x%d' % num_args),
            make_binary(num_args - 1))
    return Function(
        ['x' + str(x + 1) for x in range(num_args)],
        make_binary(num_args))


# mapping from name -> dict or Function
# the dict is a mapping from arity -> Function
EXPRESSIONS_CASE = {
    'SafeDiv': {
        2: Function(['n', 'd'], Case(
            [(EQ(NamedAttributeRef('d'), NumericLiteral(0)),
              NumericLiteral(0))],
            DIVIDE(NamedAttributeRef('n'), NamedAttributeRef('d')))),
        3: Function(['n', 'd', 'default'], Case(
            [(EQ(NamedAttributeRef('d'), NumericLiteral(0)),
              NamedAttributeRef('default'))],
            DIVIDE(NamedAttributeRef('n'), NamedAttributeRef('d'))))
    },
    'TheAnswerToLifeTheUniverseAndEverything': Function(
        [], NumericLiteral(42)),
    'greatest': lambda num_args: create_nested_binary(num_args, GREATER),
    'least': lambda num_args: create_nested_binary(num_args, LESSER),
    'greater': create_nested_binary(2, GREATER),
    'lesser': create_nested_binary(2, LESSER),
    'substr': Function(['str', 'begin', 'end'],
                       SUBSTR([NamedAttributeRef('str'),
                               NamedAttributeRef('begin'),
                               NamedAttributeRef('end')
                               ])),
    'len': Function(['str'], LEN(NamedAttributeRef('str'))),
    'head': Function(['str', 'length'], Case(
        [(GT(NamedAttributeRef('length'), LEN(NamedAttributeRef('str'))),
         NamedAttributeRef('str'))],
        SUBSTR([NamedAttributeRef('str'),
                NumericLiteral(0),
                NamedAttributeRef('length')
                ]))),
    'tail': Function(['str', 'length'], Case(
        [(GT(NamedAttributeRef('length'), LEN(NamedAttributeRef('str'))),
         NamedAttributeRef('str'))],
        SUBSTR([NamedAttributeRef('str'),
                MINUS(LEN(NamedAttributeRef('str')),
                      NamedAttributeRef('length')
                      ),
                PLUS(LEN(NamedAttributeRef('str')), NumericLiteral(1))
                ])))
}<commit_msg>use LESSER and GREATER to simplify the head and tail function<commit_after>'''
A library of expressions that can be composed of existing expressions.
'''

from .udf import Function
from raco.expression import *


def lookup(function_name, num_args):
    func = EXPRESSIONS.get(function_name.lower())
    if hasattr(func, '__call__'):
        return func(num_args)
    if isinstance(func, dict):
        return func.get(num_args)
    return func


def create_nested_binary(num_args, func):
    if num_args < 2:
        return None

    def make_binary(num_args):
        if num_args == 2:
            return func(
                NamedAttributeRef('x%d' % num_args),
                NamedAttributeRef('x%d' % (num_args - 1)))
        return func(
            NamedAttributeRef('x%d' % num_args),
            make_binary(num_args - 1))
    return Function(
        ['x' + str(x + 1) for x in range(num_args)],
        make_binary(num_args))


# mapping from name -> dict or Function
# the dict is a mapping from arity -> Function
EXPRESSIONS_CASE = {
    'SafeDiv': {
        2: Function(['n', 'd'], Case(
            [(EQ(NamedAttributeRef('d'), NumericLiteral(0)),
              NumericLiteral(0))],
            DIVIDE(NamedAttributeRef('n'), NamedAttributeRef('d')))),
        3: Function(['n', 'd', 'default'], Case(
            [(EQ(NamedAttributeRef('d'), NumericLiteral(0)),
              NamedAttributeRef('default'))],
            DIVIDE(NamedAttributeRef('n'), NamedAttributeRef('d'))))
    },
    'TheAnswerToLifeTheUniverseAndEverything': Function(
        [], NumericLiteral(42)),
    'greatest': lambda num_args: create_nested_binary(num_args, GREATER),
    'least': lambda num_args: create_nested_binary(num_args, LESSER),
    'greater': create_nested_binary(2, GREATER),
    'lesser': create_nested_binary(2, LESSER),
    'substr': Function(['str', 'begin', 'end'],
                       SUBSTR([NamedAttributeRef('str'),
                               NamedAttributeRef('begin'),
                               NamedAttributeRef('end')
                               ])),
    'len': Function(['str'], LEN(NamedAttributeRef('str'))),
    'head': Function(['str', 'length'],
                     SUBSTR([NamedAttributeRef('str'),
                             NumericLiteral(0),
                             LESSER(LEN(NamedAttributeRef('str')),
                                    NamedAttributeRef('length'))
                             ])),
    'tail': Function(['str', 'length'],
                     SUBSTR([NamedAttributeRef('str'),
                             GREATER(MINUS(LEN(NamedAttributeRef('str')),
                                           NamedAttributeRef('length')),
                                     NumericLiteral(0)),
                             LEN(NamedAttributeRef('str'))
                             ]))
}
{'fraction_non_alphanumeric': 0.14053958737436079, 'fraction_numerical': 0.00564274378416505, 'mean_word_length': 2.2392918332381497, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1352548', 'n_tokens_mistral': 1638, 'n_tokens_neox': 1537, 'n_words': 319}