diff --git a/extract_parametrize.py b/extract_parametrize.py index e69de29b..e094b28d 100644 --- a/extract_parametrize.py +++ b/extract_parametrize.py @@ -0,0 +1,41 @@ +import ast + +def extract_all_parametrize(filepath): + with open(filepath, "r") as f: + tree = ast.parse(f.read(), filename=filepath) + + parametrize_data = {} + + for node in ast.walk(tree): + if isinstance(node, ast.FunctionDef): + for deco in node.decorator_list: + if ( + isinstance(deco, ast.Call) and + isinstance(deco.func, ast.Attribute) and + deco.func.attr == "parametrize" + ): + names_node = deco.args[0] + values_node = deco.args[1] + + if isinstance(names_node, ast.Str): + param_names = [x.strip() for x in names_node.s.split(",")] + elif isinstance(names_node, ast.List): + param_names = [elt.s for elt in names_node.elts if isinstance(elt, ast.Str)] + else: + param_names = [] + + param_values = [] + if isinstance(values_node, ast.List): + for elt in values_node.elts: + if isinstance(elt, ast.Tuple): + value = [ast.literal_eval(e) for e in elt.elts] + else: + value = [ast.literal_eval(elt)] + param_values.append(value) + + parametrize_data[node.name] = { + "names": param_names, + "values": param_values + } + + return parametrize_data \ No newline at end of file