Skip to content

Commit b64e218

Browse files
cguardiacguardia
authored
fix: make fixup script consistent with migration docs (googleapis#208)
* Generate code consistent with the upgrade documentation. * Add an option to generate the code using keyword arguments instead of a request. * Generate stylistically consistent code (no spaces in keywords, double quotes for strings). * Reformat the script itself to use the same code styling.
1 parent 9968ca9 commit b64e218

2 files changed

Lines changed: 214 additions & 67 deletions

File tree

‎UPGRADING.md‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,9 @@ python3 -m pip install google-cloud-pubsub
3232

3333
* The script `fixup_pubsub_v1_keywords.py` is shipped with the library. It expects
3434
an input directory (with the code to convert) and an empty destination directory.
35+
Optionally, the `--use-keywords` switch can be added to generate flattened keyword
36+
parameters instead of a request dictionary (see the following section for an
37+
explanation).
3538

3639
```sh
3740
$ scripts/fixup_pubsub_v1_keywords.py --input-directory .samples/ --output-directory samples/

‎scripts/fixup_pubsub_v1_keywords.py‎

Lines changed: 211 additions & 67 deletions
Original file line numberDiff line numberDiff line change
@@ -20,12 +20,11 @@
2020
import libcst as cst
2121
import pathlib
2222
import sys
23-
from typing import (Any, Callable, Dict, List, Sequence, Tuple)
23+
from typing import Any, Callable, Dict, List, Sequence, Tuple
2424

2525

2626
def partition(
27-
predicate: Callable[[Any], bool],
28-
iterator: Sequence[Any]
27+
predicate: Callable[[Any], bool], iterator: Sequence[Any]
2928
) -> Tuple[List[Any], List[Any]]:
3029
"""A stable, out-of-place partition."""
3130
results = ([], [])
@@ -38,40 +37,128 @@ def partition(
3837

3938

4039
class pubsubCallTransformer(cst.CSTTransformer):
41-
CTRL_PARAMS: Tuple[str] = ('retry', 'timeout', 'metadata')
40+
CTRL_PARAMS: Tuple[str] = ("retry", "timeout", "metadata")
4241
METHOD_TO_PARAMS: Dict[str, Tuple[str]] = {
43-
'acknowledge': ('subscription', 'ack_ids', ),
44-
'create_snapshot': ('name', 'subscription', 'labels', ),
45-
'create_subscription': ('name', 'topic', 'push_config', 'ack_deadline_seconds', 'retain_acked_messages', 'message_retention_duration', 'labels', 'enable_message_ordering', 'expiration_policy', 'filter', 'dead_letter_policy', 'retry_policy', 'detached', ),
46-
'create_topic': ('name', 'labels', 'message_storage_policy', 'kms_key_name', ),
47-
'delete_snapshot': ('snapshot', ),
48-
'delete_subscription': ('subscription', ),
49-
'delete_topic': ('topic', ),
50-
'detach_subscription': ('subscription', ),
51-
'get_snapshot': ('snapshot', ),
52-
'get_subscription': ('subscription', ),
53-
'get_topic': ('topic', ),
54-
'list_snapshots': ('project', 'page_size', 'page_token', ),
55-
'list_subscriptions': ('project', 'page_size', 'page_token', ),
56-
'list_topics': ('project', 'page_size', 'page_token', ),
57-
'list_topic_snapshots': ('topic', 'page_size', 'page_token', ),
58-
'list_topic_subscriptions': ('topic', 'page_size', 'page_token', ),
59-
'modify_ack_deadline': ('subscription', 'ack_ids', 'ack_deadline_seconds', ),
60-
'modify_push_config': ('subscription', 'push_config', ),
61-
'publish': ('topic', 'messages', ),
62-
'pull': ('subscription', 'max_messages', 'return_immediately', ),
63-
'seek': ('subscription', 'time', 'snapshot', ),
64-
'streaming_pull': ('subscription', 'stream_ack_deadline_seconds', 'ack_ids', 'modify_deadline_seconds', 'modify_deadline_ack_ids', 'client_id', 'max_outstanding_messages', 'max_outstanding_bytes', ),
65-
'update_snapshot': ('snapshot', 'update_mask', ),
66-
'update_subscription': ('subscription', 'update_mask', ),
67-
'update_topic': ('topic', 'update_mask', ),
68-
69-
'get_iam_policy': ('resource', 'options', ),
70-
'set_iam_policy': ('resource', 'policy', ),
71-
'test_iam_permissions': ('resource', 'permissions', ),
72-
42+
"acknowledge": (
43+
"subscription",
44+
"ack_ids",
45+
),
46+
"create_snapshot": (
47+
"name",
48+
"subscription",
49+
"labels",
50+
),
51+
"create_subscription": (
52+
"name",
53+
"topic",
54+
"push_config",
55+
"ack_deadline_seconds",
56+
"retain_acked_messages",
57+
"message_retention_duration",
58+
"labels",
59+
"enable_message_ordering",
60+
"expiration_policy",
61+
"filter",
62+
"dead_letter_policy",
63+
"retry_policy",
64+
"detached",
65+
),
66+
"create_topic": (
67+
"name",
68+
"labels",
69+
"message_storage_policy",
70+
"kms_key_name",
71+
),
72+
"delete_snapshot": ("snapshot",),
73+
"delete_subscription": ("subscription",),
74+
"delete_topic": ("topic",),
75+
"detach_subscription": ("subscription",),
76+
"get_snapshot": ("snapshot",),
77+
"get_subscription": ("subscription",),
78+
"get_topic": ("topic",),
79+
"list_snapshots": (
80+
"project",
81+
"page_size",
82+
"page_token",
83+
),
84+
"list_subscriptions": (
85+
"project",
86+
"page_size",
87+
"page_token",
88+
),
89+
"list_topics": (
90+
"project",
91+
"page_size",
92+
"page_token",
93+
),
94+
"list_topic_snapshots": (
95+
"topic",
96+
"page_size",
97+
"page_token",
98+
),
99+
"list_topic_subscriptions": (
100+
"topic",
101+
"page_size",
102+
"page_token",
103+
),
104+
"modify_ack_deadline": (
105+
"subscription",
106+
"ack_ids",
107+
"ack_deadline_seconds",
108+
),
109+
"modify_push_config": (
110+
"subscription",
111+
"push_config",
112+
),
113+
"pull": (
114+
"subscription",
115+
"max_messages",
116+
"return_immediately",
117+
),
118+
"seek": (
119+
"subscription",
120+
"time",
121+
"snapshot",
122+
),
123+
"streaming_pull": (
124+
"subscription",
125+
"stream_ack_deadline_seconds",
126+
"ack_ids",
127+
"modify_deadline_seconds",
128+
"modify_deadline_ack_ids",
129+
"client_id",
130+
"max_outstanding_messages",
131+
"max_outstanding_bytes",
132+
),
133+
"update_snapshot": (
134+
"snapshot",
135+
"update_mask",
136+
),
137+
"update_subscription": (
138+
"subscription",
139+
"update_mask",
140+
),
141+
"update_topic": (
142+
"topic",
143+
"update_mask",
144+
),
145+
"get_iam_policy": (
146+
"resource",
147+
"options",
148+
),
149+
"set_iam_policy": (
150+
"resource",
151+
"policy",
152+
),
153+
"test_iam_permissions": (
154+
"resource",
155+
"permissions",
156+
),
73157
}
74158

159+
def __init__(self, use_keywords=False):
160+
self._use_keywords = use_keywords
161+
75162
def leave_Call(self, original: cst.Call, updated: cst.Call) -> cst.CSTNode:
76163
try:
77164
key = original.func.attr.value
@@ -88,35 +175,80 @@ def leave_Call(self, original: cst.Call, updated: cst.Call) -> cst.CSTNode:
88175
return updated
89176

90177
kwargs, ctrl_kwargs = partition(
91-
lambda a: not a.keyword.value in self.CTRL_PARAMS,
92-
kwargs
178+
lambda a: not a.keyword.value in self.CTRL_PARAMS, kwargs
93179
)
94180

95-
args, ctrl_args = args[:len(kword_params)], args[len(kword_params):]
96-
ctrl_kwargs.extend(cst.Arg(value=a.value, keyword=cst.Name(value=ctrl))
97-
for a, ctrl in zip(ctrl_args, self.CTRL_PARAMS))
181+
args, ctrl_args = args[: len(kword_params)], args[len(kword_params) :]
182+
ctrl_kwargs.extend(
183+
cst.Arg(
184+
value=a.value,
185+
keyword=cst.Name(value=ctrl),
186+
equal=cst.AssignEqual(
187+
whitespace_before=cst.SimpleWhitespace(""),
188+
whitespace_after=cst.SimpleWhitespace(""),
189+
),
190+
)
191+
for a, ctrl in zip(ctrl_args, self.CTRL_PARAMS)
192+
)
98193

99-
request_arg = cst.Arg(
100-
value=cst.Dict([
101-
cst.DictElement(
102-
cst.SimpleString("'{}'".format(name)),
103-
cst.Element(value=arg.value)
194+
if self._use_keywords:
195+
new_kwargs = [
196+
cst.Arg(
197+
value=arg.value,
198+
keyword=cst.Name(value=name),
199+
equal=cst.AssignEqual(
200+
whitespace_before=cst.SimpleWhitespace(""),
201+
whitespace_after=cst.SimpleWhitespace(""),
202+
),
104203
)
105-
# Note: the args + kwargs looks silly, but keep in mind that
106-
# the control parameters had to be stripped out, and that
107-
# those could have been passed positionally or by keyword.
108-
for name, arg in zip(kword_params, args + kwargs)]),
109-
keyword=cst.Name("request")
110-
)
204+
for name, arg in zip(kword_params, args + kwargs)
205+
]
206+
new_kwargs.extend(
207+
[
208+
cst.Arg(
209+
value=arg.value,
210+
keyword=cst.Name(value=arg.keyword.value),
211+
equal=cst.AssignEqual(
212+
whitespace_before=cst.SimpleWhitespace(""),
213+
whitespace_after=cst.SimpleWhitespace(""),
214+
),
215+
)
216+
for arg in ctrl_kwargs
217+
]
218+
)
219+
return updated.with_changes(args=new_kwargs)
220+
else:
221+
request_arg = cst.Arg(
222+
value=cst.Dict(
223+
[
224+
cst.DictElement(
225+
cst.SimpleString('"{}"'.format(name)),
226+
cst.Element(value=arg.value),
227+
)
228+
for name, arg in zip(kword_params, args + kwargs)
229+
]
230+
+ [
231+
cst.DictElement(
232+
cst.SimpleString('"{}"'.format(arg.keyword.value)),
233+
cst.Element(value=arg.value),
234+
)
235+
for arg in ctrl_kwargs
236+
]
237+
),
238+
keyword=cst.Name("request"),
239+
equal=cst.AssignEqual(
240+
whitespace_before=cst.SimpleWhitespace(""),
241+
whitespace_after=cst.SimpleWhitespace(""),
242+
),
243+
)
111244

112-
return updated.with_changes(
113-
args=[request_arg] + ctrl_kwargs
114-
)
245+
return updated.with_changes(args=[request_arg])
115246

116247

117248
def fix_files(
118249
in_dir: pathlib.Path,
119250
out_dir: pathlib.Path,
251+
use_keywords: bool = False,
120252
*,
121253
transformer=pubsubCallTransformer(),
122254
):
@@ -129,11 +261,12 @@ def fix_files(
129261
pyfile_gen = (
130262
pathlib.Path(os.path.join(root, f))
131263
for root, _, files in os.walk(in_dir)
132-
for f in files if os.path.splitext(f)[1] == ".py"
264+
for f in files
265+
if os.path.splitext(f)[1] == ".py"
133266
)
134267

135268
for fpath in pyfile_gen:
136-
with open(fpath, 'r') as f:
269+
with open(fpath, "r") as f:
137270
src = f.read()
138271

139272
# Parse the code and insert method call fixes.
@@ -145,11 +278,11 @@ def fix_files(
145278
updated_path.parent.mkdir(parents=True, exist_ok=True)
146279

147280
# Generate the updated source file at the corresponding path.
148-
with open(updated_path, 'w') as f:
281+
with open(updated_path, "w") as f:
149282
f.write(updated.code)
150283

151284

152-
if __name__ == '__main__':
285+
if __name__ == "__main__":
153286
parser = argparse.ArgumentParser(
154287
description="""Fix up source that uses the pubsub client library.
155288
@@ -164,24 +297,34 @@ def fix_files(
164297
165298
These all constitute false negatives. The tool will also detect false
166299
positives when an API method shares a name with another method.
167-
""")
300+
"""
301+
)
168302
parser.add_argument(
169-
'-d',
170-
'--input-directory',
303+
"-d",
304+
"--input-directory",
171305
required=True,
172-
dest='input_dir',
173-
help='the input directory to walk for python files to fix up',
306+
dest="input_dir",
307+
help="the input directory to walk for python files to fix up",
174308
)
175309
parser.add_argument(
176-
'-o',
177-
'--output-directory',
310+
"-o",
311+
"--output-directory",
178312
required=True,
179-
dest='output_dir',
180-
help='the directory to output files fixed via un-flattening',
313+
dest="output_dir",
314+
help="the directory to output files fixed via un-flattening",
315+
)
316+
parser.add_argument(
317+
"-k",
318+
"--use-keywords",
319+
required=False,
320+
action="store_true",
321+
dest="use_keywords",
322+
help="Use keyword arguments instead of constructing a request",
181323
)
182324
args = parser.parse_args()
183325
input_dir = pathlib.Path(args.input_dir)
184326
output_dir = pathlib.Path(args.output_dir)
327+
use_keywords = args.use_keywords
185328
if not input_dir.is_dir():
186329
print(
187330
f"input directory '{input_dir}' does not exist or is not a directory",
@@ -203,4 +346,5 @@ def fix_files(
203346
)
204347
sys.exit(-1)
205348

206-
fix_files(input_dir, output_dir)
349+
transformer = pubsubCallTransformer(use_keywords=use_keywords)
350+
fix_files(input_dir, output_dir, use_keywords, transformer=transformer)

0 commit comments

Comments
 (0)