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 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}")