def test_nested_import_from(): """ Tests whether the WIR Extraction works for nested from imports """ test_code = cleandoc(""" from mlinspect.utils import get_project_root print(get_project_root()) """) test_ast = ast.parse(test_code) extractor = WirExtractor(test_ast) extracted_wir = extractor.extract_wir() expected_graph = networkx.DiGraph() expected_import = WirNode(0, "mlinspect.utils", "Import", CodeReference(1, 0, 1, 44)) expected_call_one = WirNode(1, "get_project_root", "Call", CodeReference(3, 6, 3, 24)) expected_graph.add_edge(expected_import, expected_call_one, type="caller", arg_index=-1) expected_call_two = WirNode(2, "print", "Call", CodeReference(3, 0, 3, 25)) expected_graph.add_edge(expected_call_one, expected_call_two, type="input", arg_index=0) compare(networkx.to_dict_of_dicts(extracted_wir), networkx.to_dict_of_dicts(expected_graph))
def test_string_call_attribute(): """ Tests whether the WIR Extraction works for a very simple attribute call """ test_code = cleandoc(""" "hello ".join("world") """) test_ast = ast.parse(test_code) extractor = WirExtractor(test_ast) extracted_wir = extractor.extract_wir() expected_graph = networkx.DiGraph() expected_constant_one = WirNode(0, "hello ", "Constant", CodeReference(1, 0, 1, 8)) expected_constant_two = WirNode(1, "world", "Constant", CodeReference(1, 14, 1, 21)) expected_attribute_call = WirNode(2, "join", "Call", CodeReference(1, 0, 1, 22)) expected_graph.add_edge(expected_constant_one, expected_attribute_call, type="caller", arg_index=-1) expected_graph.add_edge(expected_constant_two, expected_attribute_call, type="input", arg_index=0) compare(networkx.to_dict_of_dicts(extracted_wir), networkx.to_dict_of_dicts(expected_graph))
def test_list_creation(): """ Tests whether the WIR Extraction works for lists """ test_code = cleandoc(""" print(["test1", "test2"]) """) test_ast = ast.parse(test_code) extractor = WirExtractor(test_ast) extracted_wir = extractor.extract_wir() expected_graph = networkx.DiGraph() expected_constant_one = WirNode(0, "test1", "Constant", CodeReference(1, 7, 1, 14)) expected_constant_two = WirNode(1, "test2", "Constant", CodeReference(1, 16, 1, 23)) expected_list = WirNode(2, "as_list", "List", CodeReference(1, 6, 1, 24)) expected_graph.add_edge(expected_constant_one, expected_list, type="input", arg_index=0) expected_graph.add_edge(expected_constant_two, expected_list, type="input", arg_index=1) expected_call = WirNode(3, "print", "Call", CodeReference(1, 0, 1, 25)) expected_graph.add_edge(expected_list, expected_call, type="input", arg_index=0) compare(networkx.to_dict_of_dicts(extracted_wir), networkx.to_dict_of_dicts(expected_graph))
def test_print_expressions(): """ Tests whether the WIR Extraction works for an expression with very simple nested calls """ test_code = cleandoc(""" print("test".isupper()) """) test_ast = ast.parse(test_code) extractor = WirExtractor(test_ast) extracted_wir = extractor.extract_wir() expected_graph = networkx.DiGraph() expected_constant = WirNode(0, "test", "Constant", CodeReference(1, 6, 1, 12)) expected_call_one = WirNode(1, "isupper", "Call", CodeReference(1, 6, 1, 22)) expected_graph.add_edge(expected_constant, expected_call_one, type="caller", arg_index=-1) expected_call_two = WirNode(2, "print", "Call", CodeReference(1, 0, 1, 23)) expected_graph.add_edge(expected_call_one, expected_call_two, type="input", arg_index=0) compare(networkx.to_dict_of_dicts(extracted_wir), networkx.to_dict_of_dicts(expected_graph))
def test_import_from(): """ Tests whether the WIR Extraction works for from imports """ test_code = cleandoc(""" from math import sqrt sqrt(4) """) test_ast = ast.parse(test_code) extractor = WirExtractor(test_ast) extracted_wir = extractor.extract_wir() expected_graph = networkx.DiGraph() expected_import = WirNode(0, "math", "Import", CodeReference(1, 0, 1, 21)) expected_constant = WirNode(1, "4", "Constant", CodeReference(3, 5, 3, 6)) expected_constant_call = WirNode(2, "sqrt", "Call", CodeReference(3, 0, 3, 7)) expected_graph.add_edge(expected_import, expected_constant_call, type="caller", arg_index=-1) expected_graph.add_edge(expected_constant, expected_constant_call, type="input", arg_index=0) compare(networkx.to_dict_of_dicts(extracted_wir), networkx.to_dict_of_dicts(expected_graph))
def test_call_after_call(): """ Tests whether the WIR Extraction works for a very simple attribute call """ test_code = cleandoc(""" "hello ".capitalize().count() """) test_ast = ast.parse(test_code) extractor = WirExtractor(test_ast) extracted_wir = extractor.extract_wir() expected_graph = networkx.DiGraph() expected_constant_one = WirNode(0, "hello ", "Constant", CodeReference(1, 0, 1, 8)) expected_call_one = WirNode(1, "capitalize", "Call", CodeReference(1, 0, 1, 21)) expected_call_two = WirNode(2, "count", "Call", CodeReference(1, 0, 1, 29)) expected_graph.add_edge(expected_constant_one, expected_call_one, type="caller", arg_index=-1) expected_graph.add_edge(expected_call_one, expected_call_two, type="caller", arg_index=-1) compare(networkx.to_dict_of_dicts(extracted_wir), networkx.to_dict_of_dicts(expected_graph))
def test_print_var_usage(): """ Tests whether the WIR Extraction works for a very simple var usage """ test_code = cleandoc(""" test_var = "test" print(test_var)""") test_ast = ast.parse(test_code) extractor = WirExtractor(test_ast) extracted_wir = extractor.extract_wir() expected_graph = networkx.DiGraph() expected_constant = WirNode(0, "test", "Constant", CodeReference(1, 11, 1, 17)) expected_assign = WirNode(1, "test_var", "Assign", CodeReference(1, 0, 1, 17)) expected_graph.add_edge(expected_constant, expected_assign, type="input", arg_index=0) expected_call = WirNode(2, "print", "Call", CodeReference(2, 0, 2, 15)) expected_graph.add_node(expected_call) expected_graph.add_edge(expected_assign, expected_call, type="input", arg_index=0) compare(networkx.to_dict_of_dicts(extracted_wir), networkx.to_dict_of_dicts(expected_graph))
def test_index_subscript_with_module_information(): """ Tests whether the WIR Extraction works for lists """ test_code = cleandoc(""" import pandas as pd data = pd.read_csv('test_path') data['income-per-year'] """) test_ast = ast.parse(test_code) extractor = WirExtractor(test_ast) module_info = { CodeReference(3, 7, 3, 31): ('pandas.io.parsers', 'read_csv'), CodeReference(4, 0, 4, 23): ('pandas.core.frame', '__getitem__') } extracted_wir = extractor.extract_wir() extractor.add_runtime_info(module_info, {}) expected_graph = networkx.DiGraph() expected_import = WirNode(0, "pandas", "Import", CodeReference(1, 0, 1, 19)) expected_constant_one = WirNode(1, "test_path", "Constant", CodeReference(3, 19, 3, 30)) expected_call = WirNode(2, "read_csv", "Call", CodeReference(3, 7, 3, 31), ('pandas.io.parsers', 'read_csv')) expected_graph.add_edge(expected_import, expected_call, type="caller", arg_index=-1) expected_graph.add_edge(expected_constant_one, expected_call, type="input", arg_index=0) expected_assign = WirNode(3, "data", "Assign", CodeReference(3, 0, 3, 31)) expected_graph.add_edge(expected_call, expected_assign, type="input", arg_index=0) expected_constant_two = WirNode(4, "income-per-year", "Constant", CodeReference(4, 5, 4, 22)) expected_index_subscript = WirNode(5, "Index-Subscript", "Subscript", CodeReference(4, 0, 4, 23), ('pandas.core.frame', '__getitem__')) expected_graph.add_edge(expected_assign, expected_index_subscript, type="caller", arg_index=-1) expected_graph.add_edge(expected_constant_two, expected_index_subscript, type="input", arg_index=0) compare(networkx.to_dict_of_dicts(extracted_wir), networkx.to_dict_of_dicts(expected_graph))
def test_keyword(): """ Tests whether the WIR Extraction works for function calls with keyword usage """ test_code = cleandoc(""" print('comma', 'separated', 'words', sep=', ') """) test_ast = ast.parse(test_code) extractor = WirExtractor(test_ast) extracted_wir = extractor.extract_wir() expected_graph = networkx.DiGraph() expected_constant_one = WirNode(0, "comma", "Constant", CodeReference(1, 6, 1, 13)) expected_constant_two = WirNode(1, "separated", "Constant", CodeReference(1, 15, 1, 26)) expected_constant_three = WirNode(2, "words", "Constant", CodeReference(1, 28, 1, 35)) expected_constant_four = WirNode(3, ", ", "Constant", CodeReference(1, 41, 1, 45)) expected_keyword = WirNode(4, "sep", "Keyword") expected_call = WirNode(5, "print", "Call", CodeReference(1, 0, 1, 46)) expected_graph.add_edge(expected_constant_four, expected_keyword, type="input", arg_index=0) expected_graph.add_edge(expected_constant_one, expected_call, type="input", arg_index=0) expected_graph.add_edge(expected_constant_two, expected_call, type="input", arg_index=1) expected_graph.add_edge(expected_constant_three, expected_call, type="input", arg_index=2) expected_graph.add_edge(expected_keyword, expected_call, type="input", arg_index=3) compare(networkx.to_dict_of_dicts(extracted_wir), networkx.to_dict_of_dicts(expected_graph))
def test_pipeline_executor_function_call_info_extraction(): """ Tests whether the capturing of module information works """ test_code = get_pandas_read_csv_and_dropna_code() pipeline_executor.singleton = pipeline_executor.PipelineExecutor() pipeline_executor.singleton.run(None, None, test_code, []) expected_module_info = { CodeReference(5, 13, 5, 85): ('posixpath', 'join'), CodeReference(5, 26, 5, 49): ('builtins', 'str'), CodeReference(5, 30, 5, 48): ('mlinspect.utils', 'get_project_root'), CodeReference(6, 11, 6, 34): ('pandas.io.parsers', 'read_csv'), CodeReference(7, 7, 7, 24): ('pandas.core.frame', 'dropna') } compare(pipeline_executor.singleton.code_reference_to_module, expected_module_info)
def test_print_stmt(): """ Tests whether the WIR Extraction works for a very simple print statement """ test_code = "print('test')" test_ast = ast.parse(test_code) extractor = WirExtractor(test_ast) extracted_wir = extractor.extract_wir() expected_graph = networkx.DiGraph() expected_constant = WirNode(0, "test", "Constant", CodeReference(1, 6, 1, 12)) expected_call = WirNode(1, "print", "Call", CodeReference(1, 0, 1, 13)) expected_graph.add_edge(expected_constant, expected_call, type="input", arg_index=0) compare(networkx.to_dict_of_dicts(extracted_wir), networkx.to_dict_of_dicts(expected_graph))
def test_pipeline_executor_function_subscript_index_info_extraction(): """ Tests whether the capturing of module information works """ test_code = cleandoc(""" import os import pandas as pd from mlinspect.utils import get_project_root train_file = os.path.join(str(get_project_root()), "test", "data", "adult_train.csv") raw_data = pd.read_csv(train_file, na_values='?', index_col=0) data = raw_data.dropna() data['income-per-year'] """) pipeline_executor.singleton = pipeline_executor.PipelineExecutor() pipeline_executor.singleton.run(None, None, test_code, []) expected_module_info = { CodeReference(5, 13, 5, 85): ('posixpath', 'join'), CodeReference(5, 26, 5, 49): ('builtins', 'str'), CodeReference(5, 30, 5, 48): ('mlinspect.utils', 'get_project_root'), CodeReference(6, 11, 6, 62): ('pandas.io.parsers', 'read_csv'), CodeReference(7, 7, 7, 24): ('pandas.core.frame', 'dropna'), CodeReference(8, 0, 8, 23): ('pandas.core.frame', '__getitem__') } compare(pipeline_executor.singleton.code_reference_to_module, expected_module_info)
def test_index_assign(): """ Tests whether the WIR Extraction works for lists """ test_code = cleandoc(""" import pandas as pd data = pd.read_csv('test_path') data['label'] = "test" """) test_ast = ast.parse(test_code) extractor = WirExtractor(test_ast) extracted_wir = extractor.extract_wir() expected_graph = networkx.DiGraph() expected_import = WirNode(0, "pandas", "Import", CodeReference(1, 0, 1, 19)) expected_constant_one = WirNode(1, "test_path", "Constant", CodeReference(3, 19, 3, 30)) expected_call = WirNode(2, "read_csv", "Call", CodeReference(3, 7, 3, 31)) expected_graph.add_edge(expected_import, expected_call, type="caller", arg_index=-1) expected_graph.add_edge(expected_constant_one, expected_call, type="input", arg_index=0) expected_assign = WirNode(3, "data", "Assign", CodeReference(3, 0, 3, 31)) expected_graph.add_edge(expected_call, expected_assign, type="input", arg_index=0) expected_constant_two = WirNode(4, "label", "Constant", CodeReference(4, 5, 4, 12)) expected_graph.add_node(expected_constant_two) expected_constant_three = WirNode(5, "test", "Constant", CodeReference(4, 16, 4, 22)) expected_graph.add_node(expected_constant_three) expected_subscript_assign = WirNode(6, 'data.label', 'Subscript-Assign', CodeReference(4, 0, 4, 13)) expected_graph.add_edge(expected_assign, expected_subscript_assign, type="caller", arg_index=-1) compare(networkx.to_dict_of_dicts(extracted_wir), networkx.to_dict_of_dicts(expected_graph))
def get_call_description_info(): """ Get the module info for the adult_easy pipeline """ call_description_info = { CodeReference(lineno=12, col_offset=11, end_lineno=12, end_col_offset=62): 'adult_train.csv', CodeReference(lineno=14, col_offset=7, end_lineno=14, end_col_offset=24): 'dropna', CodeReference(lineno=16, col_offset=38, end_lineno=16, end_col_offset=61): "to ['income-per-year']", CodeReference(lineno=16, col_offset=9, end_lineno=16, end_col_offset=89): "label_binarize, classes: ['>50K', '<=50K']", CodeReference(lineno=19, col_offset=20, end_lineno=19, end_col_offset=72): 'Categorical Encoder (OneHotEncoder)', CodeReference(lineno=20, col_offset=16, end_lineno=20, end_col_offset=46): 'Numerical Encoder (StandardScaler)', CodeReference(lineno=26, col_offset=19, end_lineno=26, end_col_offset=48): 'Decision Tree' } return call_description_info
def test_tuple_assign(): """ Tests whether the WIR Extraction works for a very simple var usage """ test_code = cleandoc(""" x, y = (1, 2) print(x)""") test_ast = ast.parse(test_code) extractor = WirExtractor(test_ast) extracted_wir = extractor.extract_wir() expected_graph = networkx.DiGraph() expected_constant_one = WirNode(1, '1', 'Constant', CodeReference(1, 8, 1, 9)) expected_constant_two = WirNode(2, '2', 'Constant', CodeReference(1, 11, 1, 12)) expetected_constant_tuple = WirNode(3, 'as_tuple', 'Tuple', CodeReference(1, 7, 1, 13)) expected_graph.add_edge(expected_constant_one, expetected_constant_tuple, type="input", arg_index=0) expected_graph.add_edge(expected_constant_two, expetected_constant_tuple, type="input", arg_index=1) expected_var_x = WirNode(4, 'x', 'Assign', CodeReference(1, 0, 1, 1)) expected_var_y = WirNode(5, 'y', 'Assign', CodeReference(1, 3, 1, 4)) expected_graph.add_edge(expetected_constant_tuple, expected_var_x, type="input", arg_index=0) expected_graph.add_edge(expetected_constant_tuple, expected_var_y, type="input", arg_index=0) expected_call = WirNode(6, 'print', 'Call', CodeReference(2, 0, 2, 8)) expected_graph.add_edge(expected_var_x, expected_call, type="input", arg_index=0) compare(networkx.to_dict_of_dicts(extracted_wir), networkx.to_dict_of_dicts(expected_graph))
def get_module_info(): """ Get the module info for the adult_easy pipeline """ module_info = { CodeReference(lineno=10, col_offset=0, end_lineno=10, end_col_offset=23): ('builtins', 'print'), CodeReference(lineno=11, col_offset=30, end_lineno=11, end_col_offset=48): ('mlinspect.utils', 'get_project_root'), CodeReference(lineno=11, col_offset=26, end_lineno=11, end_col_offset=49): ('builtins', 'str'), CodeReference(lineno=11, col_offset=13, end_lineno=11, end_col_offset=85): ('posixpath', 'join'), CodeReference(lineno=12, col_offset=11, end_lineno=12, end_col_offset=62): ('pandas.io.parsers', 'read_csv'), CodeReference(lineno=14, col_offset=7, end_lineno=14, end_col_offset=24): ('pandas.core.frame', 'dropna'), CodeReference(lineno=16, col_offset=38, end_lineno=16, end_col_offset=61): ('pandas.core.frame', '__getitem__', 'Projection'), CodeReference(lineno=16, col_offset=9, end_lineno=16, end_col_offset=89): ('sklearn.preprocessing._label', 'label_binarize'), CodeReference(lineno=19, col_offset=20, end_lineno=19, end_col_offset=72): ('sklearn.preprocessing._encoders', 'OneHotEncoder'), CodeReference(lineno=20, col_offset=16, end_lineno=20, end_col_offset=46): ('sklearn.preprocessing._data', 'StandardScaler'), CodeReference(lineno=18, col_offset=25, end_lineno=21, end_col_offset=2): ('sklearn.compose._column_transformer', 'ColumnTransformer'), CodeReference(lineno=26, col_offset=19, end_lineno=26, end_col_offset=48): ('sklearn.tree._classes', 'DecisionTreeClassifier'), CodeReference(lineno=24, col_offset=18, end_lineno=26, end_col_offset=51): ('sklearn.pipeline', 'Pipeline'), CodeReference(lineno=28, col_offset=0, end_lineno=28, end_col_offset=33): ('sklearn.pipeline', 'fit'), CodeReference(lineno=31, col_offset=0, end_lineno=31, end_col_offset=26): ('builtins', 'print') } return module_info
def get_expected_dag_adult_easy_py(): """ Get the expected DAG for the adult_easy pipeline """ # pylint: disable=too-many-locals expected_graph = networkx.DiGraph() expected_data_source = DagNode(18, OperatorType.DATA_SOURCE, CodeReference(12, 11, 12, 62), ('pandas.io.parsers', 'read_csv'), "adult_train.csv") expected_graph.add_node(expected_data_source) expected_select = DagNode(20, OperatorType.SELECTION, CodeReference(14, 7, 14, 24), ('pandas.core.frame', 'dropna'), "dropna") expected_graph.add_edge(expected_data_source, expected_select) expected_train_data = DagNode(56, OperatorType.TRAIN_DATA, CodeReference(24, 18, 26, 51), ('sklearn.pipeline', 'fit', 'Train Data')) expected_graph.add_edge(expected_select, expected_train_data) expected_pipeline_project_one = DagNode( 34, OperatorType.PROJECTION, CodeReference(18, 25, 21, 2), ('sklearn.compose._column_transformer', 'ColumnTransformer', 'Projection'), "to ['education'] (ColumnTransformer)") expected_graph.add_edge(expected_train_data, expected_pipeline_project_one) expected_pipeline_project_two = DagNode( 35, OperatorType.PROJECTION, CodeReference(18, 25, 21, 2), ('sklearn.compose._column_transformer', 'ColumnTransformer', 'Projection'), "to ['workclass'] (ColumnTransformer)") expected_graph.add_edge(expected_train_data, expected_pipeline_project_two) expected_pipeline_project_three = DagNode( 40, OperatorType.PROJECTION, CodeReference(18, 25, 21, 2), ('sklearn.compose._column_transformer', 'ColumnTransformer', 'Projection'), "to ['age'] (ColumnTransformer)") expected_graph.add_edge(expected_train_data, expected_pipeline_project_three) expected_pipeline_project_four = DagNode( 41, OperatorType.PROJECTION, CodeReference(18, 25, 21, 2), ('sklearn.compose._column_transformer', 'ColumnTransformer', 'Projection'), "to ['hours-per-week'] (ColumnTransformer)") expected_graph.add_edge(expected_train_data, expected_pipeline_project_four) expected_pipeline_transformer_one = DagNode( 34, OperatorType.TRANSFORMER, CodeReference(19, 20, 19, 72), ('sklearn.preprocessing._encoders', 'OneHotEncoder', 'Pipeline'), "Categorical Encoder (OneHotEncoder), Column: 'education'") expected_graph.add_edge(expected_pipeline_project_one, expected_pipeline_transformer_one) expected_pipeline_transformer_two = DagNode( 35, OperatorType.TRANSFORMER, CodeReference(19, 20, 19, 72), ('sklearn.preprocessing._encoders', 'OneHotEncoder', 'Pipeline'), "Categorical Encoder (OneHotEncoder), Column: 'workclass'") expected_graph.add_edge(expected_pipeline_project_two, expected_pipeline_transformer_two) expected_pipeline_transformer_three = DagNode( 40, OperatorType.TRANSFORMER, CodeReference(20, 16, 20, 46), ('sklearn.preprocessing._data', 'StandardScaler', 'Pipeline'), "Numerical Encoder (StandardScaler), Column: 'age'") expected_graph.add_edge(expected_pipeline_project_three, expected_pipeline_transformer_three) expected_pipeline_transformer_four = DagNode( 41, OperatorType.TRANSFORMER, CodeReference(20, 16, 20, 46), ('sklearn.preprocessing._data', 'StandardScaler', 'Pipeline'), "Numerical Encoder (StandardScaler), Column: 'hours-per-week'") expected_graph.add_edge(expected_pipeline_project_four, expected_pipeline_transformer_four) expected_pipeline_concatenation = DagNode( 46, OperatorType.CONCATENATION, CodeReference(18, 25, 21, 2), ('sklearn.compose._column_transformer', 'ColumnTransformer', 'Concatenation')) expected_graph.add_edge(expected_pipeline_transformer_one, expected_pipeline_concatenation) expected_graph.add_edge(expected_pipeline_transformer_two, expected_pipeline_concatenation) expected_graph.add_edge(expected_pipeline_transformer_three, expected_pipeline_concatenation) expected_graph.add_edge(expected_pipeline_transformer_four, expected_pipeline_concatenation) expected_estimator = DagNode( 51, OperatorType.ESTIMATOR, CodeReference(26, 19, 26, 48), ('sklearn.tree._classes', 'DecisionTreeClassifier', 'Pipeline'), "Decision Tree") expected_graph.add_edge(expected_pipeline_concatenation, expected_estimator) expected_pipeline_fit = DagNode(56, OperatorType.FIT, CodeReference(24, 18, 26, 51), ('sklearn.pipeline', 'fit', 'Pipeline')) expected_graph.add_edge(expected_estimator, expected_pipeline_fit) expected_project = DagNode( 23, OperatorType.PROJECTION, CodeReference(16, 38, 16, 61), ('pandas.core.frame', '__getitem__', 'Projection'), "to ['income-per-year']") expected_graph.add_edge(expected_select, expected_project) expected_project_modify = DagNode( 28, OperatorType.PROJECTION_MODIFY, CodeReference(16, 9, 16, 89), ('sklearn.preprocessing._label', 'label_binarize'), "label_binarize, classes: ['>50K', '<=50K']") expected_graph.add_edge(expected_project, expected_project_modify) expected_train_labels = DagNode( 56, OperatorType.TRAIN_LABELS, CodeReference(24, 18, 26, 51), ('sklearn.pipeline', 'fit', 'Train Labels')) expected_graph.add_edge(expected_project_modify, expected_train_labels) expected_graph.add_edge(expected_train_labels, expected_pipeline_fit) return expected_graph
def test_tuples(): """ Tests whether the WIR Extraction works for tuples """ test_code = cleandoc(""" from sklearn import preprocessing ('categorical', preprocessing.OneHotEncoder(handle_unknown='ignore'), ['education', 'workclass']) """) test_ast = ast.parse(test_code) extractor = WirExtractor(test_ast) extracted_wir = extractor.extract_wir() expected_graph = networkx.DiGraph() expected_import_from = WirNode(0, "sklearn", "Import", CodeReference(1, 0, 1, 33)) expected_constant_one = WirNode(1, "categorical", "Constant", CodeReference(3, 1, 3, 14)) expected_constant_two = WirNode(2, "ignore", "Constant", CodeReference(3, 59, 3, 67)) expected_keyword = WirNode(3, "handle_unknown", "Keyword") expected_graph.add_edge(expected_constant_two, expected_keyword, type="input", arg_index=0) expected_call = WirNode(4, "OneHotEncoder", "Call", CodeReference(3, 16, 3, 68)) expected_graph.add_edge(expected_import_from, expected_call, type="caller", arg_index=-1) expected_graph.add_edge(expected_keyword, expected_call, type="input", arg_index=0) expected_constant_three = WirNode(5, "education", "Constant", CodeReference(3, 71, 3, 82)) expected_constant_four = WirNode(6, "workclass", "Constant", CodeReference(3, 84, 3, 95)) expected_list = WirNode(7, "as_list", "List", CodeReference(3, 70, 3, 96)) expected_graph.add_edge(expected_constant_three, expected_list, type="input", arg_index=0) expected_graph.add_edge(expected_constant_four, expected_list, type="input", arg_index=1) expected_tuple = WirNode(8, "as_tuple", "Tuple", CodeReference(3, 0, 3, 97)) expected_graph.add_edge(expected_constant_one, expected_tuple, type="input", arg_index=0) expected_graph.add_edge(expected_call, expected_tuple, type="input", arg_index=1) expected_graph.add_edge(expected_list, expected_tuple, type="input", arg_index=2) compare(networkx.to_dict_of_dicts(extracted_wir), networkx.to_dict_of_dicts(expected_graph))
def get_expected_result(): """ Get the expected PrintFirstRowsAnalyzer(2) result for the adult_easy example """ expected_result = { DagNode(node_id=18, operator_type=OperatorType.DATA_SOURCE, module=('pandas.io.parsers', 'read_csv'), code_reference=CodeReference(12, 11, 12, 62), description='adult_train.csv'): [ InspectionInputRow( values=[46, 'Private', 128645, 'Some-college', 10, 'Divorced', 'Prof-specialty', 'Not-in-family', 'White', 'Female', 0, 0, 40, 'United-States', '<=50K'], fields=['age', 'workclass', 'fnlwgt', 'education', 'education-num', 'marital-status', 'occupation', 'relationship', 'race', 'sex', 'capital-gain', 'capital-loss', 'hours-per-week', 'native-country', 'income-per-year']), InspectionInputRow( values=[29, 'Local-gov', 115585, 'Some-college', 10, 'Never-married', 'Handlers-cleaners', 'Not-in-family', 'White', 'Male', 0, 0, 50, 'United-States', '<=50K'], fields=['age', 'workclass', 'fnlwgt', 'education', 'education-num', 'marital-status', 'occupation', 'relationship', 'race', 'sex', 'capital-gain', 'capital-loss', 'hours-per-week', 'native-country', 'income-per-year'])], DagNode(node_id=20, operator_type=OperatorType.SELECTION, module=('pandas.core.frame', 'dropna'), code_reference=CodeReference(14, 7, 14, 24), description='dropna'): [ InspectionInputRow( values=[46, 'Private', 128645, 'Some-college', 10, 'Divorced', 'Prof-specialty', 'Not-in-family', 'White', 'Female', 0, 0, 40, 'United-States', '<=50K'], fields=['age', 'workclass', 'fnlwgt', 'education', 'education-num', 'marital-status', 'occupation', 'relationship', 'race', 'sex', 'capital-gain', 'capital-loss', 'hours-per-week', 'native-country', 'income-per-year']), InspectionInputRow( values=[29, 'Local-gov', 115585, 'Some-college', 10, 'Never-married', 'Handlers-cleaners', 'Not-in-family', 'White', 'Male', 0, 0, 50, 'United-States', '<=50K'], fields=['age', 'workclass', 'fnlwgt', 'education', 'education-num', 'marital-status', 'occupation', 'relationship', 'race', 'sex', 'capital-gain', 'capital-loss', 'hours-per-week', 'native-country', 'income-per-year'])], DagNode(node_id=23, operator_type=OperatorType.PROJECTION, module=('pandas.core.frame', '__getitem__', 'Projection'), code_reference=CodeReference(16, 38, 16, 61), description="to ['income-per-year']"): [ InspectionInputRow(values=['<=50K'], fields=['array']), InspectionInputRow(values=['<=50K'], fields=['array'])], DagNode(node_id=28, operator_type=OperatorType.PROJECTION_MODIFY, module=('sklearn.preprocessing._label', 'label_binarize'), code_reference=CodeReference(16, 9, 16, 89), description="label_binarize, classes: ['>50K', '<=50K']"): [ InspectionInputRow(values=[array(1)], fields=['array']), InspectionInputRow(values=[array(1)], fields=['array'])], DagNode(node_id=56, operator_type=OperatorType.TRAIN_DATA, module=('sklearn.pipeline', 'fit', 'Train Data'), code_reference=CodeReference(24, 18, 26, 51), description=None): [ InspectionInputRow( values=[46, 'Private', 128645, 'Some-college', 10, 'Divorced', 'Prof-specialty', 'Not-in-family', 'White', 'Female', 0, 0, 40, 'United-States', '<=50K', 1], fields=['age', 'workclass', 'fnlwgt', 'education', 'education-num', 'marital-status', 'occupation', 'relationship', 'race', 'sex', 'capital-gain', 'capital-loss', 'hours-per-week', 'native-country', 'income-per-year', 'mlinspect_index']), InspectionInputRow( values=[29, 'Local-gov', 115585, 'Some-college', 10, 'Never-married', 'Handlers-cleaners', 'Not-in-family', 'White', 'Male', 0, 0, 50, 'United-States', '<=50K', 2], fields=['age', 'workclass', 'fnlwgt', 'education', 'education-num', 'marital-status', 'occupation', 'relationship', 'race', 'sex', 'capital-gain', 'capital-loss', 'hours-per-week', 'native-country', 'income-per-year', 'mlinspect_index'])], DagNode(node_id=56, operator_type=OperatorType.TRAIN_LABELS, module=('sklearn.pipeline', 'fit', 'Train Labels'), code_reference=CodeReference(24, 18, 26, 51), description=None): [ InspectionInputRow(values=[array(1)], fields=['array']), InspectionInputRow(values=[array(1)], fields=['array'])], DagNode(node_id=40, operator_type=OperatorType.PROJECTION, code_reference=CodeReference(18, 25, 21, 2), module=('sklearn.compose._column_transformer', 'ColumnTransformer', 'Projection'), description="to ['age'] (ColumnTransformer)"): [ InspectionInputRow(values=[46], fields=['age']), InspectionInputRow(values=[29], fields=['age'])], DagNode(node_id=34, operator_type=OperatorType.PROJECTION, code_reference=CodeReference(18, 25, 21, 2), module=('sklearn.compose._column_transformer', 'ColumnTransformer', 'Projection'), description="to ['education'] (ColumnTransformer)"): [ InspectionInputRow(values=['Some-college'], fields=['education']), InspectionInputRow(values=['Some-college'], fields=['education'])], DagNode(node_id=41, operator_type=OperatorType.PROJECTION, code_reference=CodeReference(18, 25, 21, 2), module=('sklearn.compose._column_transformer', 'ColumnTransformer', 'Projection'), description="to ['hours-per-week'] (ColumnTransformer)"): [ InspectionInputRow(values=[40], fields=['hours-per-week']), InspectionInputRow(values=[50], fields=['hours-per-week'])], DagNode(node_id=35, operator_type=OperatorType.PROJECTION, code_reference=CodeReference(18, 25, 21, 2), module=('sklearn.compose._column_transformer', 'ColumnTransformer', 'Projection'), description="to ['workclass'] (ColumnTransformer)"): [ InspectionInputRow(values=['Private'], fields=['workclass']), InspectionInputRow(values=['Local-gov'], fields=['workclass'])], DagNode(node_id=40, operator_type=OperatorType.TRANSFORMER, code_reference=CodeReference(20, 16, 20, 46), module=('sklearn.preprocessing._data', 'StandardScaler', 'Pipeline'), description="Numerical Encoder (StandardScaler), Column: 'age'"): [ InspectionInputRow(values=[array(RangeComparison(0.5, 0.6))], fields=['array']), InspectionInputRow(values=[array(RangeComparison(-0.8, -0.7))], fields=['array'])], DagNode(node_id=41, operator_type=OperatorType.TRANSFORMER, code_reference=CodeReference(20, 16, 20, 46), module=('sklearn.preprocessing._data', 'StandardScaler', 'Pipeline'), description="Numerical Encoder (StandardScaler), Column: 'hours-per-week'"): [ InspectionInputRow(values=[array(RangeComparison(-0.09, -0.08))], fields=['array']), InspectionInputRow(values=[array(RangeComparison(0.7, 0.8))], fields=['array'])], DagNode(node_id=34, operator_type=OperatorType.TRANSFORMER, code_reference=CodeReference(19, 20, 19, 72), module=('sklearn.preprocessing._encoders', 'OneHotEncoder', 'Pipeline'), description="Categorical Encoder (OneHotEncoder), Column: 'education'"): [ InspectionInputRow( values=[array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1.])], fields=['array']), InspectionInputRow( values=[array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1.])], fields=['array'])], DagNode(node_id=35, operator_type=OperatorType.TRANSFORMER, code_reference=CodeReference(19, 20, 19, 72), module=('sklearn.preprocessing._encoders', 'OneHotEncoder', 'Pipeline'), description="Categorical Encoder (OneHotEncoder), Column: 'workclass'"): [ InspectionInputRow(values=[array([0., 0., 1., 0., 0., 0., 0.])], fields=['array']), InspectionInputRow(values=[array([0., 1., 0., 0., 0., 0., 0.])], fields=['array'])], DagNode(node_id=46, operator_type=OperatorType.CONCATENATION, code_reference=CodeReference(18, 25, 21, 2), module=('sklearn.compose._column_transformer', 'ColumnTransformer', 'Concatenation'), description=None): [ InspectionInputRow(values=[array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 1., 0., 0., 0., 0., RangeComparison(0.5, 0.6), RangeComparison(-0.09, -0.08)])], fields=['array']), InspectionInputRow(values=[array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 1., 0., 0., 0., 0., 0., RangeComparison(-0.8, -0.7), RangeComparison(0.7, 0.8)])], fields=['array'])], DagNode(node_id=51, operator_type=OperatorType.ESTIMATOR, code_reference=CodeReference(26, 19, 26, 48), module=('sklearn.tree._classes', 'DecisionTreeClassifier', 'Pipeline'), description='Decision Tree'): None } return expected_result
def get_expected_cleaned_wir_adult_easy(): """ Get the expected cleaned WIR for the adult_easy pipeline """ # pylint: disable=too-many-locals expected_graph = networkx.DiGraph() expected_print_one = WirNode(6, "print", "Call", CodeReference(10, 0, 10, 23), ('builtins', 'print')) expected_graph.add_node(expected_print_one) expected_get_project_root = WirNode( 7, "get_project_root", "Call", CodeReference(11, 30, 11, 48), ('mlinspect.utils', 'get_project_root')) expected_str = WirNode(8, "str", "Call", CodeReference(11, 26, 11, 49), ('builtins', 'str')) expected_graph.add_edge(expected_get_project_root, expected_str) expected_join = WirNode(12, "join", "Call", CodeReference(11, 13, 11, 85), ('posixpath', 'join')) expected_graph.add_edge(expected_str, expected_join) expected_read_csv = WirNode(18, "read_csv", "Call", CodeReference(12, 11, 12, 62), ('pandas.io.parsers', 'read_csv')) expected_graph.add_edge(expected_join, expected_read_csv) expected_dropna = WirNode(20, "dropna", "Call", CodeReference(14, 7, 14, 24), ('pandas.core.frame', 'dropna')) expected_graph.add_edge(expected_read_csv, expected_dropna) expected_fit = WirNode(56, "fit", "Call", CodeReference(28, 0, 28, 33), ('sklearn.pipeline', 'fit')) expected_index_subscript = WirNode( 23, "Index-Subscript", "Subscript", CodeReference(16, 38, 16, 61), ('pandas.core.frame', '__getitem__', 'Projection')) expected_graph.add_edge(expected_dropna, expected_fit) expected_graph.add_edge(expected_dropna, expected_index_subscript) expected_label_binarize = WirNode( 28, "label_binarize", "Call", CodeReference(16, 9, 16, 89), ('sklearn.preprocessing._label', 'label_binarize')) expected_graph.add_edge(expected_index_subscript, expected_label_binarize) expected_graph.add_edge(expected_label_binarize, expected_fit) expected_one_hot_encoder = WirNode( 33, "OneHotEncoder", "Call", CodeReference(19, 20, 19, 72), ('sklearn.preprocessing._encoders', 'OneHotEncoder')) expected_standard_scaler = WirNode( 39, "StandardScaler", "Call", CodeReference(20, 16, 20, 46), ('sklearn.preprocessing._data', 'StandardScaler')) expected_column_transformer = WirNode( 46, "ColumnTransformer", "Call", CodeReference(18, 25, 21, 2), ('sklearn.compose._column_transformer', 'ColumnTransformer')) expected_graph.add_edge(expected_one_hot_encoder, expected_column_transformer) expected_graph.add_edge(expected_standard_scaler, expected_column_transformer) expected_decision_tree_classifier = WirNode( 51, "DecisionTreeClassifier", "Call", CodeReference(26, 19, 26, 48), ('sklearn.tree._classes', 'DecisionTreeClassifier')) expected_pipeline = WirNode(54, "Pipeline", "Call", CodeReference(24, 18, 26, 51), ('sklearn.pipeline', 'Pipeline')) expected_graph.add_edge(expected_column_transformer, expected_pipeline) expected_graph.add_edge(expected_decision_tree_classifier, expected_pipeline) expected_graph.add_edge(expected_pipeline, expected_fit) expected_print_two = WirNode(58, "print", "Call", CodeReference(31, 0, 31, 26), ('builtins', 'print')) expected_graph.add_node(expected_print_two) return expected_graph