def main()

in smallpond/execution/workqueue.py [0:0]


def main():
    import argparse

    from smallpond.execution.task import Probe

    parser = argparse.ArgumentParser(prog="workqueue.py", description="Work Queue Reader")
    parser.add_argument("wq_root", help="Work queue root path")
    parser.add_argument("-f", "--work_filter", default="", help="Work item filter")
    parser.add_argument("-x", "--expand_batch", action="store_true", help="Expand batched works")
    parser.add_argument("-c", "--count_object", action="store_true", help="Count number of objects")
    parser.add_argument("-n", "--top_n_class", default=20, type=int, help="Show the top n classes")
    parser.add_argument("-l", "--log_level", default="INFO", help="Logging message level")
    args = parser.parse_args()

    logger.remove()
    logger.add(
        sys.stdout,
        format=r"[{time:%Y-%m-%d %H:%M:%S.%f}] {level} {message}",
        level=args.log_level,
    )

    wq = WorkQueueOnFilesystem(args.wq_root)
    for work in wq.list_works(args.expand_batch):
        if isinstance(work, Probe):
            continue
        if args.work_filter in work.key:
            logger.info(work)
            if args.count_object:
                object_cnt = {}
                count_objects(work, object_cnt)
                sorted_counts = sorted([(v, k) for k, v in object_cnt.items()], reverse=True)
                for count, class_name in sorted_counts[: args.top_n_class]:
                    logger.info(f"  {class_name}: {count}")