Coverage for skema/model_assembly/code_types.py: 32%

1from enum import Enum, auto, unique 

2 

3import networkx as nx 

4 

5 

6@unique 

7class CodeType(Enum): 

8 ACCESSOR = auto() 

9 CALCULATION = auto() 

10 CONVERSION = auto() 

11 FILEIO = auto() 

12 HELPER = auto() 

13 LOGGING = auto() 

14 MODEL = auto() 

15 PIPELINE = auto() 

16 UNKNOWN = auto() 

17 

18 

19def build_code_type_decision_tree(): 

20 G = nx.DiGraph() 

21 G.add_node( 

22 "C0", 

23 type="condition", 

24 func=lambda d: d["num_switches"] >= 1, 

25 shape="rectangle", 

26 label="num_switches >= 1", 

27 ) 

28 G.add_node( 

29 "C1", 

30 type="condition", 

31 func=lambda d: d["max_call_depth"] <= 2, 

32 shape="rectangle", 

33 label="max_call_depth <= 2", 

34 ) 

35 G.add_node( 

36 "C2", 

37 type="condition", 

38 func=lambda d: d["num_assgs"] >= 1, 

39 shape="rectangle", 

40 label="num_assgs >= 1", 

41 ) 

42 G.add_node( 

43 "C3", 

44 type="condition", 

45 func=lambda d: d["num_math_assgs"] >= 1, 

46 shape="rectangle", 

47 label="num_math_assgs >= 1", 

48 ) 

49 G.add_node( 

50 "C4", 

51 type="condition", 

52 func=lambda d: d["num_data_changes"] >= 1, 

53 shape="rectangle", 

54 label="num_data_changes >= 1", 

55 ) 

56 G.add_node( 

57 "C5", 

58 type="condition", 

59 func=lambda d: d["num_var_access"] >= 1, 

60 shape="rectangle", 

61 label="num_var_access >= 1", 

62 ) 

63 G.add_node( 

64 "C6", 

65 type="condition", 

66 func=lambda d: d["num_math_assgs"] >= 5, 

67 shape="rectangle", 

68 label="num_math_assgs >= 5", 

69 ) 

70 

71 G.add_node("Accessor", type=CodeType.ACCESSOR, color="blue") 

72 G.add_node("Calculation", type=CodeType.CALCULATION, color="blue") 

73 G.add_node("Conversion", type=CodeType.CONVERSION, color="blue") 

74 G.add_node("File I/O", type=CodeType.FILEIO, color="blue") 

75 # G.add_node("Helper", type=CodeType.HELPER, color='blue') 

76 # G.add_node("Logging", type=CodeType.LOGGING, color='blue') 

77 G.add_node("Model", type=CodeType.MODEL, color="blue") 

78 G.add_node("Pipeline", type=CodeType.PIPELINE, color="blue") 

79 G.add_node("Unknown", type=CodeType.UNKNOWN, color="blue") 

80 

81 G.add_edge("C0", "Pipeline", type=True, color="darkgreen") 

82 G.add_edge("C0", "C1", type=False, color="red") 

83 

84 G.add_edge("C1", "C2", type=True, color="darkgreen") 

85 G.add_edge("C1", "Pipeline", type=False, color="red") 

86 

87 G.add_edge("C2", "File I/O", type=False, color="red") 

88 G.add_edge("C2", "C3", type=True, color="darkgreen") 

89 

90 G.add_edge("C3", "C4", type=True, color="darkgreen") 

91 G.add_edge("C3", "C5", type=False, color="red") 

92 

93 G.add_edge("C4", "C6", type=True, color="darkgreen") 

94 G.add_edge("C4", "Conversion", type=False, color="red") 

95 

96 G.add_edge("C5", "Accessor", type=True, color="darkgreen") 

97 G.add_edge("C5", "Unknown", type=False, color="red") 

98 

99 G.add_edge("C6", "Model", type=True, color="darkgreen") 

100 G.add_edge("C6", "Calculation", type=False, color="red") 

101 

102 return G