11import os
22import pathlib
33import shutil
4+
45from typing import List
6+ from argparse import ArgumentParser
57
68from grpc_tools import command
79
@@ -35,8 +37,9 @@ def remove_protos(rootdirpath: str):
3537 os .remove (os .path .join (root , file ))
3638
3739
38- def fix_file_contents (rootdir : str ):
40+ def fix_file_contents (rootdir , protobuf_version : str ):
3941 flake_ignore_line = "# flake8: " + "noqa" # prevent ignore the file
42+ package_path = "ydb._grpc." + protobuf_version + ".protos"
4043
4144 for dirpath , _ , fnames in os .walk (rootdir ):
4245 for fname in fnames :
@@ -47,24 +50,30 @@ def fix_file_contents(rootdir: str):
4750 content = f .read ()
4851
4952 # Fix imports
50- content = content .replace ("from protos" , "from ydb._grpc.protos" )
53+ content = content .replace ("from protos" , "from " + package_path )
5154
5255 # Add ignore style check
5356 content = content .replace ("# -*- coding: utf-8 -*-" , "# -*- coding: utf-8 -*-\n " + flake_ignore_line )
5457 f .seek (0 )
5558 f .write (content )
5659
5760
58- def generate_protobuf (src_proto_dir : str , dst_dir : str ):
61+ def generate_protobuf (src_proto_dir : str , dst_dir , protobuf_version : str ):
5962 shutil .rmtree (dst_dir , ignore_errors = True )
6063
6164 shutil .copytree (src_proto_dir , dst_dir , ignore = files_filter )
62- create_init_files (dst_dir )
6365
6466 command .build_package_protos (dst_dir )
6567 remove_protos (dst_dir )
66- fix_file_contents (dst_dir )
68+ create_init_files (dst_dir )
69+ fix_file_contents (dst_dir , protobuf_version )
6770
6871
6972if __name__ == '__main__' :
70- generate_protobuf ("ydb-api-protos" , "ydb/_grpc" )
73+ parser = ArgumentParser ()
74+ parser .add_argument ("--target-version" , default = "v4" , help = "Protobuf version" )
75+
76+ args = parser .parse_args ()
77+
78+ target_dir = "ydb/_grpc/" + args .target_version
79+ generate_protobuf ("ydb-api-protos" , target_dir , args .target_version )
0 commit comments