()
| 45 | return args |
| 46 | |
| 47 | def main(): |
| 48 | args = parse_args() |
| 49 | |
| 50 | cfg = Config.fromfile(args.config) |
| 51 | |
| 52 | if args.cfg_options is None: |
| 53 | args.cfg_options = dict() |
| 54 | if args.root is not None: |
| 55 | args.cfg_options["root"] = args.root |
| 56 | if args.workdir is not None: |
| 57 | args.cfg_options["workdir"] = args.workdir |
| 58 | if args.tag is not None: |
| 59 | args.cfg_options["tag"] = args.tag |
| 60 | cfg.merge_from_dict(args.cfg_options) |
| 61 | |
| 62 | update_data_root(cfg, root=args.root) |
| 63 | |
| 64 | exp_path = os.path.join(cfg.root, cfg.workdir, cfg.tag) |
| 65 | if args.if_remove is None: |
| 66 | args.if_remove = bool(input(f"| Arguments PRESS 'y' to REMOVE: {exp_path}? ") == 'y') |
| 67 | if args.if_remove: |
| 68 | import shutil |
| 69 | shutil.rmtree(exp_path, ignore_errors=True) |
| 70 | print(f"| Arguments Remove work_dir: {exp_path}") |
| 71 | else: |
| 72 | print(f"| Arguments Keep work_dir: {exp_path}") |
| 73 | os.makedirs(exp_path, exist_ok=True) |
| 74 | |
| 75 | downloader = DOWNLOADER.build(cfg.downloader) |
| 76 | |
| 77 | print(f"| Check Downloading {cfg.tag}...") |
| 78 | stocks = downloader.check_download() |
| 79 | print(f"| Check Downloading {cfg.tag} Done! Failed: {len(stocks)} / {len(downloader.stocks)}") |
| 80 | |
| 81 | batch_size = cfg.batch_size if cfg.batch_size < len(stocks) else 5 |
| 82 | batch_size = min(len(stocks), batch_size) |
| 83 | |
| 84 | processes = [] |
| 85 | remaining_stocks = downloader.check_download() |
| 86 | |
| 87 | while remaining_stocks: |
| 88 | batch = remaining_stocks[:batch_size] |
| 89 | remaining_stocks = remaining_stocks[batch_size:] |
| 90 | |
| 91 | process = StockDownloaderProcess(batch, downloader) |
| 92 | processes.append(process) |
| 93 | process.start() |
| 94 | |
| 95 | for process in processes: |
| 96 | process.join() |
| 97 | |
| 98 | if __name__ == '__main__': |
| 99 | main() |
no test coverage detected