binary_search_tool/test/gen_obj.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# Copyright 2020 The ChromiumOS Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Script to generate a list of object files.

0 represents a good object file.
1 represents a bad object file.
"""


import argparse
import os
import random
import sys

from binary_search_tool.test import common


def Main(argv):
    """Generates a list, the value of each element is 0 or 1.

    The number of 1s in the list is specified by bad_obj_num.
    The others are all 0s. The total number of 0s and 1s is specified by obj_num.

    Args:
      argv: argument from command line

    Returns:
      0 always.
    """
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-n",
        "--obj_num",
        dest="obj_num",
        default=common.DEFAULT_OBJECT_NUMBER,
        help=("Number of total objects."),
    )
    parser.add_argument(
        "-b",
        "--bad_obj_num",
        dest="bad_obj_num",
        default=common.DEFAULT_BAD_OBJECT_NUMBER,
        help=(
            "Number of bad objects. Must be great than or "
            "equal to zero and less than total object "
            "number."
        ),
    )
    parser.add_argument(
        "-o",
        "--obj_list",
        dest="obj_list",
        default="",
        help=(
            "List of comma seperated objects to generate. "
            "A 0 means the object is good, a 1 means the "
            "object is bad."
        ),
    )
    options = parser.parse_args(argv)

    obj_num = int(options.obj_num)
    bad_obj_num = int(options.bad_obj_num)
    bad_to_gen = int(options.bad_obj_num)
    obj_list = options.obj_list
    if not obj_list:
        obj_list = []
        for i in range(obj_num):
            if bad_to_gen > 0 and random.randint(1, obj_num) <= bad_obj_num:
                obj_list.append(1)
                bad_to_gen -= 1
            else:
                obj_list.append(0)
        while bad_to_gen > 0:
            t = random.randint(0, obj_num - 1)
            if obj_list[t] == 0:
                obj_list[t] = 1
                bad_to_gen -= 1
    else:
        obj_list = obj_list.split(",")

    if os.path.isfile(common.OBJECTS_FILE):
        os.remove(common.OBJECTS_FILE)
    if os.path.isfile(common.WORKING_SET_FILE):
        os.remove(common.WORKING_SET_FILE)

    with open(common.OBJECTS_FILE, "w", encoding="utf-8") as f:
        with open(common.WORKING_SET_FILE, "w", encoding="utf-8") as w:
            for i in obj_list:
                f.write("{0}\n".format(i))
                w.write("{0}\n".format(i))

    obj_num = len(obj_list)
    bad_obj_num = obj_list.count(1)
    print(
        "Generated {0} object files, with {1} bad ones.".format(
            obj_num, bad_obj_num
        )
    )

    return 0


if __name__ == "__main__":
    retval = Main(sys.argv[1:])
    sys.exit(retval)