vlambda博客
学习文章列表

import ast# Simple code that sets a variable var to 1 and then prints it.code = """var = 1print(var)"""# Converts code to AST. Object 'head' points to the head of the AST.head = ast.parse(code)print(head)

<_ast.Module object at ...>

print(ast.dump(head))

Module( body=[ Assign( targets=[Name(id='var', ctx=Store())], value=Num(n=1) ), Expr( value=Call(func=Name(id='print', ctx=Load()), args=[Name(id='var', ctx=Load())], keywords=[]) ) ])

print(head.body[0].value.n)

1

head.body[0].value.n = 2print(ast.dump(head))

Module( body=[ Assign( targets=[Name(id='var', ctx=Store())], value=Num(n=2) ), Expr( value=Call(func=Name(id='print', ctx=Load()), args=[Name(id='var', ctx=Load())], keywords=[]) ) ])

import astunparseprint(astunparse.unparse(head))

var = 2print(var)

import pandas as pdmi = pd.MultiIndex.from_product([[1, 2], ['a', 'b']], names=['x', 'y'])print(mi.levels[0].name)mi.levels[0].name = "new name"print(mi.levels[0].name)

x new name 

import pandas as pd mi = pd.MultiIndex.from_product([[1, 2], ['a', 'b']], names=['x', 'y']) print(mi.levels[0].name) mi = mi.set_names("new name", level=0) print(mi.levels[0].name) 

import astdef is_multi_index_rename_node(node): """ Checks if the given node represents the code: <var>.levels[<idx>].name = <val> and returns the corresponding var, idx and val if it does. """ try: if ( isinstance(node, ast.Assign) and node.targets[0].attr == "name" and node.targets[0].value.value.attr == "levels" ): var = node.targets[0].value.value.value.id idx = node.targets[0].value.slice.value.n val = node.value return True, var, idx, val except: pass return False, None, None, Nonedef get_new_multi_index_rename_node(var, idx, val): """ Returns AST node that represents the code: <var> = <var>.set_names(<val>, level=<idx>) for the given var, idx and val. """ return ast.Assign( targets=[ast.Name(id=var)], value=ast.Call( func=ast.Attribute(value=ast.Name(id=var), attr="set_names"), args=[val], keywords=[ast.keyword(arg="level", value=ast.Num(n=idx))], ), )def patch(node): """ Takes an AST rooted at the give node and patches it. """ # If it is a leaf node, then no patching needed. if not hasattr(node, "_fields"): return node # For every child of the node, modify it if needed and recursively call patch on it. for (name, field) in ast.iter_fields(node): if isinstance(field, list): for i in range(len(field)): check, var, idx, val = is_multi_index_rename_node(field[i]) if check: field[i] = get_new_multi_index_rename_node(var, idx, val) else: patch(field[i]) else: check, var, idx, val = is_multi_index_rename_node(field) if check: setattr(node, name, get_new_multi_index_rename_node(var, idx, val)) else: patch(field)

from intelli_patch import patchimport astimport astunparsecode = """import pandas as pdmi = pd.MultiIndex.from_product([[1, 2], ['a', 'b']], names=['x', 'y'])mi.levels[0].name = "new name""""head = ast.parse(code)patch(head)print(astunparse.unparse(head))

import pandas as pdmi = pd.MultiIndex.from_product([[1, 2], ['a', 'b']], names=['x', 'y'])mi = mi.set_names('new name', level=0)

from intelli_patch import patchimport astimport astunparsecode = """import pandas as pdclass C(): def f(): def g(): mi.levels[ 0 ].name = "new name"mi = pd.MultiIndex.from_product([[1, 2], ['a', 'b']], names=['x', 'y'])"""head = ast.parse(code)patch(head)print(astunparse.unparse(head))

import pandas as pdclass C(): def f(): def g(): mi = mi.set_names('new name', level=0)mi = pd.MultiIndex.from_product([[1, 2], ['a', 'b']], names=['x', 'y'])

import astdef check(node): """ Takes an AST rooted at the given node and checks if there are single character variable names. """ # If it is a leaf node, then return. if not hasattr(node, "_fields"): return # For every child of the node, check if it is a variable having single character # name and recursively call check on it. for child_node in ast.iter_child_nodes(node): if isinstance(child_node, ast.Name) and len(child_node.id) == 1: print( f"Single character name '{child_node.id}' used at line number {child_node.lineno}" ) check(child_node)

from variable_name_check import checkimport astcode = """a = 1b = aprint(b)"""head = ast.parse(code)check(head)

Single character name 'a' used at line number 2Single character name 'b' used at line number 3Single character name 'a' used at line number 3Single character name 'b' used at line number 4

import astdef check(node): """ Takes an AST rooted at the given node and checks if there are un-logged except code blocks. """ # If it is a leaf node, then return. if not hasattr(node, "_fields"): return # For every child of the node, check if it is un-logged except code block # and recursively call check on it. for child_node in ast.iter_child_nodes(node): if isinstance(child_node, ast.ExceptHandler) and not is_logging_present( child_node ): print( f"Neither 'error' nor 'exception' logging is present within the except block starting at line number {child_node.lineno}" ) check(child_node)def is_logging_present(node): """ Takes an AST rooted at the given node and checks whether there is an 'error' or 'exception' logging present in it. """ # If it is a leaf node, then return False. if not hasattr(node, "_fields"): return False # If it represents an `error` or `exception`function call then return True. if ( isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute) and node.func.attr in ["error", "exception"] ): return True # Recursively checking if logging is present in the children nodes. for child_node in ast.iter_child_nodes(node): if is_logging_present(child_node): return True return False

from unlogged_except_check import checkimport astcode="""try: passexcept ValueError: logging.error("Error occurred") try: pass except KeyError: logging.exception("Exception handled") except NameError: passtry: passexcept: logging.info("Info level logging")"""head = ast.parse(code)check(head)

Neither 'error' nor 'exception' logging is present within the except block starting at line number 14Neither 'error' nor 'exception' logging is present within the except block starting at line number 19