def main():
"""Rewrite Thrift-generated Python clients to handle recursive structs. For
more details see: https://issues.apache.org/jira/browse/THRIFT-2642.
Requires package `RedBaron`, available via pip:
$ pip install redbaron
To use:
$ thrift -gen py mapd.thrift
$ mv gen-py/mapd/ttypes.py gen-py/mapd/ttypes-backup.py
$ python fix_recursive_structs.py gen-py/mapd/ttypes-backup.py gen-py/mapd/ttypes.py
"""
in_file = open(sys.argv[1], 'r')
out_file = open(sys.argv[2], 'w')
red_ast = RedBaron(in_file.read())
thrift_specs = [ts.parent for ts in red_ast.find_all(
'name', 'thrift_spec') if ts.parent.type == 'assignment' and ts.parent.parent.name in ['TDatumVal', 'TColumnData']]
nodes = []
for ts in thrift_specs:
node = ts.copy()
node.target = ts.parent.name + '.' + str(node.target)
nodes.append(node)
ts.value = 'None'
red_ast.extend(nodes)
out_file.write(red_ast.dumps())
def test_root_as_line_proxy_list_extend():
red = RedBaron("\n\na\nb\nc\n")
red.extend(["zob"])
assert red.dumps() == "\n\na\nb\nc\nzob\n"
def test_root_as_line_proxy_list_extend():
red = RedBaron("\n\na\nb\nc\n")
red.extend(["zob"])
assert red.dumps() == "\n\na\nb\nc\nzob\n"
def create_operation_stub(file_path: str):
"""Given filepath, generates a stub file at if any function in that file is decorated with the
@operation decorator.
The decorator is only identified by name - we do not ensure that it is the specific decorator.
Any decorator called "operation" would trigger a stub generation.
"""
with open(file_path) as f:
content = f.read()
red = RedBaron(content)
operations = red(
"funcdef", # get all function definitions
lambda node: any([dec("name", "operation") for dec in node.decorators],
), # with decorator named "operation"
)
if operations:
tree = RedBaron("import typing\nimport pyinfra\n")
global_args = _get_global_arg_names()
ellipses = RedBaron("...")
for operation in operations:
stub = operation.copy()
stub.decorators = []
stub = stub.copy()
stub.value = ellipses
stub.arguments.extend(global_args)
# *args and **kwargs must be moved the end of the argument list
list_argument_index = None
for i, argument in enumerate(stub.arguments):
if argument.type == "list_argument":
list_argument_index = i
if list_argument_index is not None:
list_arg = stub.arguments[list_argument_index]
stub.arguments.pop(
list_argument_index,
) # pop does not return the element for some reason
stub.arguments.append(list_arg)
dict_argument_index = None
for i, argument in enumerate(stub.arguments):
if argument.type == "dict_argument":
dict_argument_index = i
if dict_argument_index is not None:
dict_arg = stub.arguments[dict_argument_index]
stub.arguments.pop(dict_argument_index)
stub.arguments.append(dict_arg)
tree.extend(RedBaron("\n\n"))
tree.extend(stub)
output_file_path = file_path + "i"
pyi_content = tree.dumps() + "\n\n"
# If using @operation(...) as opposed to @operation, it will generate a single long
# for each function. Formatting with black solves this.
pyi_content_formatted = black.format_str(
pyi_content, mode=black.FileMode(is_pyi=True))
with open(output_file_path, "w") as f:
f.write(pyi_content_formatted)
print(f"--> Generated: {output_file_path}")
else:
logger.info(f"No operations found in file {file_path}")
