1
1
import os
2
2
import pathlib
3
3
import shutil
4
+
4
5
from typing import List
6
+ from argparse import ArgumentParser
5
7
6
8
from grpc_tools import command
7
9
@@ -35,8 +37,9 @@ def remove_protos(rootdirpath: str):
35
37
os .remove (os .path .join (root , file ))
36
38
37
39
38
- def fix_file_contents (rootdir : str ):
40
+ def fix_file_contents (rootdir , protobuf_version : str ):
39
41
flake_ignore_line = "# flake8: " + "noqa" # prevent ignore the file
42
+ package_path = "ydb._grpc." + protobuf_version + ".protos"
40
43
41
44
for dirpath , _ , fnames in os .walk (rootdir ):
42
45
for fname in fnames :
@@ -47,24 +50,30 @@ def fix_file_contents(rootdir: str):
47
50
content = f .read ()
48
51
49
52
# Fix imports
50
- content = content .replace ("from protos" , "from ydb._grpc.protos" )
53
+ content = content .replace ("from protos" , "from " + package_path )
51
54
52
55
# Add ignore style check
53
56
content = content .replace ("# -*- coding: utf-8 -*-" , "# -*- coding: utf-8 -*-\n " + flake_ignore_line )
54
57
f .seek (0 )
55
58
f .write (content )
56
59
57
60
58
- def generate_protobuf (src_proto_dir : str , dst_dir : str ):
61
+ def generate_protobuf (src_proto_dir : str , dst_dir , protobuf_version : str ):
59
62
shutil .rmtree (dst_dir , ignore_errors = True )
60
63
61
64
shutil .copytree (src_proto_dir , dst_dir , ignore = files_filter )
62
- create_init_files (dst_dir )
63
65
64
66
command .build_package_protos (dst_dir )
65
67
remove_protos (dst_dir )
66
- fix_file_contents (dst_dir )
68
+ create_init_files (dst_dir )
69
+ fix_file_contents (dst_dir , protobuf_version )
67
70
68
71
69
72
if __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