generate_commands_from_synopsis.py

import argparse
from functools import partial
import pandas as pd
from collections import Counter, defaultdict
import numpy as np
from tqdm.auto import tqdm
import warnings
import re

import sys
sys.path.append('../clai/utils')
import bashlint
from bashlint import data_tools


def update_graph(cmd, graph):
    parsed = data_tools.bash_parser(cmd)
    child = parsed.children[0]
    if not isinstance(child, bashlint.nast.PipelineNode):
        return
    prev_name = ""
    for c in child.children:
        if c.is_utility():
            cur_name = c.value
        else:
            cur_name = ""
        if prev_name and cur_name:
            graph[prev_name].add(cur_name)
        prev_name = cur_name

def add_utilities(cmd, counter):
    def get_utilities_fun(node):
        utilities = []
        if node.is_utility():
            utilities.append(node.value)
            for child in node.children:
                utilities.extend(get_utilities_fun(child))
        elif not node.is_argument():
            for child in node.children:
                utilities.extend(get_utilities_fun(child))
        return utilities

    parsed = data_tools.bash_parser(cmd)
    utils = get_utilities_fun(parsed)
    counter.update(utils)

def number_of_required_arguments(cmd):
    return sum([not x.optional for x in bashlint.grammar.bg.grammar[cmd].positional_arguments if isinstance(x, bashlint.grammar.ArgumentState)])


def get_options(name, manpage, alias_to_idx):
    change_name = {
        'grep': 'egrep',
        'gcc': 'aarch64-linux-gnu-gcc-8',
        'vim': 'rvim',
        'rename': 'file-rename'
    }

    name = change_name.get(name, name)

    x = manpage[manpage['name'] == name]
    if len(x) < 1:
        x = alias_to_idx.get(name, None)
        if x is None or len(x) < 1:
            return None
        else:
            x = manpage.loc[x]
    options = []
    paragraphs = x.iloc[0]['paragraphs']
    for p in paragraphs:
        text = p['text'].strip()
        if not p['is_option']:
            super_option_regex = r"^(?:\<[^>]+\>)?(-{1,2}\w+)(?:\<\/[^>]+\>)?[\t ]*"

            match = re.match(super_option_regex, text)
            if not match:
                continue
            p['short'] = []
            p['long'] = []
            p['expectsarg'] = None
            p['argument'] = None
            while match:
                found_text = match.group(1)
                if found_text.startswith("--"):
                    p['long'].append(found_text)
                else:
                    p['short'].append(found_text)

                text = re.sub(super_option_regex, "", text, 1)

                match = re.match(super_option_regex, text)
        if 'short' not in p:
            continue
        option = {
            'short': p['short'],
            'long': p['long'],
            'expectsarg': p['expectsarg'],
            'argument': p['argument']
        }
        if '\n' in text:
            text = text[text.find('\n') + 1:].strip()
        while text.startswith('<'):
            if '\n' in text:
                text = text[text.find('\n') + 1:].strip()
            else:
                text = re.sub("^\<[^>]+\>.*?\<\/[^>]+\>\s*", "", text).strip()
        option['text'] = text
        options.append(option)
    return options

def generate_single_command(results, avg_options=3, idx=None):
    if idx is None:
        idx = np.random.randint(len(results))
    row = results.iloc[idx]
    options = row['options']
    cmd = [row['cmd']]

    text = [row['synopsis']]

    if options:
        p = np.random.rand(len(options))
        if len(options) > avg_options:
            p = p < avg_options / len(options)
        else:
            p = p < 0.3
        p_options = [x for x, pp in zip(options, p) if pp]

        for x in p_options:
            if '--help' in x['long'] or '--version' in x['long'] or '-help' in x['short']:
                if len(p_options) > 1:
                    continue
            option_variants = x['short'] + x['long']
            option_var = np.random.choice(option_variants)
            current_option = [option_var]
            if x['expectsarg']:
                current_option.append('ARG')

            add_text = " ".join(x['text'].lower().split()[:5])
            text.append(add_text)
            cmd.append(tuple(current_option))
        cmd[1:] = sorted(cmd[1:], key=lambda x: "2"+x[0][2:] if x[0].startswith('--') else "1"+x[0][1:])
        real_cmd = [cmd[0]]
        for x in cmd[1:]:
            real_cmd.extend(x)
        cmd = real_cmd

    cmd.extend(["ARG"] * row['required'])
    return " ".join(cmd), " ".join(text)

def get_cmd_name(cmd):
    space_idx = cmd.find(' ')
    if space_idx == -1:
        return cmd
    return cmd[:space_idx]


def generate_commands(results, graph, avg_options=3, pipe_prob=0.3):
    idx = None
    cmds = []
    texts = []
    while True:
        cmd, text = generate_single_command(results, avg_options=avg_options, idx=idx)
        cmds.append(cmd)
        texts.append(text)
        p = np.random.rand()

        if p > pipe_prob:
            break

        cmd_name = get_cmd_name(cmd)
        next = graph.get(cmd_name, None)
        if next is None:
            break
        next_cmd = np.random.choice(list(graph[cmd_name]))
        idx = np.where(results['cmd'] == next_cmd)[0][0]

    cmd = " | ".join(cmds)
    text = " and ".join(texts)

    extra_cmds = ["{} | wc -l", "{} | grep ARG", "VAR=$({})", "VAR=`{}`"]
    extra_texts = [
        ["How many {}", "Count lines of {}", "Get number of {}"],
        ["Find all {} with ARG", "Which {} has ARG"],
        ["Set variable VAR to the {}", "Set variable VAR to {}"],
        ["Set variable VAR to the {}", "Set variable VAR to {}"]
    ]
    if np.random.rand() < 0.01:
        cmd_idx = np.random.randint(len(extra_cmds))
        cmd = extra_cmds[cmd_idx].format(cmd)
        text = np.random.choice(extra_texts[cmd_idx]).format(text)

    return cmd, text


def main(args):
    nl2bash = pd.read_json(args.nl2bash).T

    graph = defaultdict(lambda: set())
    with warnings.catch_warnings():
        warnings.simplefilter("ignore", FutureWarning)
        tqdm.pandas(desc="Extracting utilities graph")
    nl2bash['cmd'].progress_apply(partial(update_graph, graph=graph))

    count_utilities = Counter()
    with warnings.catch_warnings():
        warnings.simplefilter("ignore", FutureWarning)
        tqdm.pandas(desc="Extracting utilities from examples")
    nl2bash['cmd'].progress_apply(partial(add_utilities, counter=count_utilities))
    all_commands = list(bashlint.grammar.bg.grammar.keys())
    count_utilities.update(all_commands)

    commands = pd.DataFrame.from_dict(count_utilities, orient='index', columns=["count"]).reset_index() \
        .rename(columns={'index': 'cmd'}).sort_values('count').reset_index(drop=True)
    commands['required'] = commands['cmd'].apply(number_of_required_arguments)
    print(f"Found {len(commands)} total utilities")

    manpage = pd.read_json(args.manpage, lines=True)

    commands = commands.merge(manpage[['name', 'synopsis']], left_on='cmd', right_on='name', how='left')
    commands.loc[commands['synopsis'].isna(), 'synopsis'] = ''

    alias_to_idx = defaultdict(lambda: [])

    def get_aliases(x):
        idx = x.name
        for y in x['aliases']:
            y = y[0]
            alias_to_idx[y].append(idx)

    manpage.apply(get_aliases, axis=1)

    commands.drop_duplicates(inplace=True)
    print(f"Now {len(commands)} utilities")

    commands['options'] = commands['cmd'].apply(partial(get_options, manpage=manpage, alias_to_idx=alias_to_idx))
    del manpage
    results = []
    for t in tqdm(range(args.size), desc="Generating examples"):
        results.append(list(generate_commands(commands, graph)))
    results = pd.DataFrame(results, columns=["cmd", "query"])
    results.to_csv(args.output, index=False)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="Crawl examples from manpage-data.json")
    parser.add_argument("nl2bash", type=str, help="path to nl2bash-data.json")
    parser.add_argument("manpage", type=str, help="path to manpage-data.json")
    parser.add_argument("--size", type=int, default=10000, help="number of generated examples")
    parser.add_argument("-o", "--output", type=str, default="generated-examples.csv", help="path to output file")
    args = parser.parse_args()
    main(args)