helper to do the boring stuff of https://github.com/asciimoo/searx/issues/302

transform_engines.py

#!/usr/bin/env python3
#
# helper to do the boring stuff of https://github.com/asciimoo/searx/issues/302
# usage : ```python utils/transform_engines.py```
#
# use redbaron (pip install redbaron)
# have a look to : https://libcst.readthedocs.io/en/latest/why_libcst.html
#
# perhaps some engine specific QA could be done using redbaron, libcst, ast or similar libraries.

from os import listdir
from os.path import isfile, join
import redbaron


# the output will contains features = { .... "key": value, ... }
def kv_features_id(key, value):
    return 'features', key, str(value)


# the output will contains metadata = { .... "key": value, ... }
def kv_metadata_id(key, value):
    return 'metadata', key, str(value)


def kv_using_api(key, value):
    if value is not None:
        value = value.strip()
        if value.lower().startswith('\"no'):
            value = 'False'
        elif value.lower().startswith('\"yes'):
            value = 'True'
    return 'metadata', 'use_api', value


# list of global variables to map to something else
# value of the map : lambda(key,value): (map_name, new_key, new_value_as_string)
global_map = {
    'categories': kv_features_id,
    'paging': kv_features_id,
    # remove the _support suffix
    'language_support': lambda k,v : ('features', 'language', str(v)),
    # remove the _support suffix
    'time_range_support': lambda k,v : ('features', 'time_range', str(v)),
    'safesearch': kv_features_id,
}

# comments with @website, @using-api, ...
comment_map = {
    'website': kv_metadata_id,
    'using-api': kv_using_api,
}


def dict_to_python(name, dictionary):
    result = name + ' = {\n'
    for k, value in dictionary.items():
        if value is not None:
            result +=  '    "' + str(k) + '": ' + str(value) + "\n"
    result += '}'
    return result


def insert_after_with_comment(node, comment, code):
    node.insert_after(code)
    node.insert_after('# ' + comment)
    node.insert_after('\n')


def set_comment_value(new_values, line):
    for k,m in comment_map.items():
        if line.startswith('@' + k):
            v = '"' + line[len(k) + 1:].strip() + '"'
            d, k, v = m(k,v)
            new_values[d][k] = v + ','
            return True
    return False


def update_red(red):
    # get the data for the new nodes, get the nodes to remove
    new_values = {
        'features': {},
        'metadata': {}
    }
    insert_code_after = 0
    nodes_to_remove = []
    for i, node in enumerate(red):
        if isinstance(node, redbaron.StringNode) and i == 0:
            # parse string aka multiline comment
            for line in node.value.splitlines():
                set_comment_value(new_values, line.strip())
        elif isinstance(node, redbaron.CommentNode):
            # parse comment
            if 'engine dependent config' in node.value.strip().lower():
                nodes_to_remove.append(node)
            else:
                set_comment_value(new_values, node.value[1:].strip())
        elif isinstance(node, (redbaron.ImportNode, redbaron.FromImportNode, redbaron.TryNode)):
            # insert the new variables after last import including the pattern "try: import ... except: ..."
            insert_code_after = i
        elif isinstance(node, redbaron.AssignmentNode):
            # parse global variable
            name = str(node.target)
            if name in global_map:
                nodes_to_remove.append(node)
                d, k, v = global_map[name](name, node.value)
                v = v + ','
                # add trailing comments (anything actually) until the end of the line.
                while True:
                    node = node.next
                    if isinstance(node, redbaron.EndlNode):
                        break
                    v += node.dumps()
                    nodes_to_remove.append(node)
                new_values[d][k] = v
            if name == 'logger':
                # insert the new variables after "logger"
                insert_code_after = i

    # remove nodes
    for node in nodes_to_remove:
        red.remove(node)
    nodes_to_remove = []

    # remove double endl
    first_endl = None
    endl_count = 0
    for i, node in enumerate(red):
        if isinstance(node, redbaron.EndlNode):
            endl_count += 1
            if endl_count == 1:
                first_endl = node
        elif isinstance(node, redbaron.StringNode):
            # comment
            first_endl = None
            endl_count = 0
        elif isinstance(node, redbaron.CommentNode):
            pass
        else:
            if isinstance(node, redbaron.DefNode):
                if endl_count>2:
                    nodes_to_remove.append(first_endl)
                # stop on first function
                break
            elif endl_count>1:
                nodes_to_remove.append(first_endl)
            first_endl = None
            endl_count = 0

    # remove nodes
    for node in nodes_to_remove:
        red.remove(node)
    nodes_to_remove = []

    # default values
    new_values['features'].setdefault('categories', "['general'],")
    new_values['features'].setdefault('paging', 'False,')
    new_values['features'].setdefault('language', 'False,')
    new_values['features'].setdefault('time_range', 'False,')
    new_values['features'].setdefault('safesearch', 'False,')
    new_values['features'].setdefault('allow_http', 'False,')
    new_values['features'].setdefault('session_required', 'False,')
    new_values['features'].setdefault('multiple_requests', 'False,')
    # new_values['features'].setdefault('special_search_syntax', 'None,')

    new_values['metadata'].setdefault('website', '"FIXME",')
    new_values['metadata'].setdefault('use_api', 'False,')
    new_values['metadata'].setdefault('require_api_key', 'False,')
    # new_values['metadata'].setdefault('query_example', 'None,')

    # insert new nodes
    insert_after_with_comment(red[insert_code_after], 'metadata', dict_to_python('metadata', new_values['metadata']))
    insert_after_with_comment(red[insert_code_after], 'features', dict_to_python('features', new_values['features']))


def parse_content(content):
    red = redbaron.RedBaron(content)
    update_red(red)
    return red.dumps()


def parse_file(filename):
    with open(filename) as f:
        content = f.read()
    new_content = parse_content(content)
    with open(filename, 'w') as f:
        f.write(new_content)


def engine_file_names():
    engines_path = 'searx/engines'
    for f in listdir(engines_path):
        filename = join(engines_path, f)
        if isfile(filename) and filename.endswith('.py') and f != '__init__.py':
            yield filename


for f in engine_file_names():
    print(f)
    parse_file(f)