Coverage for skema/program_analysis/python2cast.py: 63%

1import os 

2import sys 

3import ast 

4import json 

5import argparse 

6from pathlib import Path 

7from tempfile import TemporaryDirectory 

8 

9from skema.program_analysis.CAST.pythonAST import py_ast_to_cast 

10from skema.program_analysis.CAST2FN import cast 

11from skema.program_analysis.CAST2FN.model.cast import SourceRef 

12from skema.program_analysis.CAST2FN.cast import CAST 

13from skema.program_analysis.CAST2FN.visitors.cast_to_agraph_visitor import ( 

14 CASTToAGraphVisitor, 

15) 

16 

17from skema.program_analysis.CAST.python.ts2cast import TS2CAST 

18from skema.program_analysis.python_preprocessor import preprocess 

19from typing import Optional 

20 

21def get_args(): 

22 parser = argparse.ArgumentParser( 

23 "Runs Python to CAST pipeline on input Python source file." 

24 ) 

25 parser.add_argument( 

26 "--rawjson", 

27 help="Dumps out raw JSON contents to stdout", 

28 action="store_true", 

29 ) 

30 parser.add_argument( 

31 "--stdout", 

32 help="Dumps CAST JSON to stdout instead of a file", 

33 action="store_true", 

34 ) 

35 parser.add_argument( 

36 "--agraph", 

37 help="Generates visualization of CAST as a PDF file", 

38 action="store_true", 

39 ) 

40 parser.add_argument( 

41 "--legacy", 

42 help="Generate CAST for GrFN 2.2 pipeline", 

43 action="store_true", 

44 ) 

45 parser.add_argument( 

46 "--ts", 

47 help="Generate CAST using tree-sitter parser generator instead of the Python AST", 

48 action="store_true", 

49 ) 

50 parser.add_argument("pyfile_path", help="input Python source file") 

51 options = parser.parse_args() 

52 return options 

53 

54 

55def python_to_cast( 

56 pyfile_path, 

57 agraph=False, 

58 astprint=False, 

59 std_out=False, 

60 rawjson=False, 

61 legacy=False, 

62 cast_obj=False, 

63 tree_sitter=False 

64) -> Optional[CAST]: 

65 """Create a CAST object from a Python file and serialize it to JSON. 

66 

67 Args: 

68 pyfile_path: Path to the Python source file 

69 agraph: If true, a PDF visualization of the graph is created. 

70 std_out: If true, the CAST JSON is printed to stdout instead 

71 of written to a file. 

72 rawjson: If true, the raw JSON contents are printed to stdout. 

73 legacy: If true, generate CAST for GrFN 2.2 pipeline. 

74 cast_obj: If true, returns the CAST as an object instead of printing to 

75 stdout. 

76 

77 Returns: 

78 If cast_obj is set to True, returns the CAST as an object. Else, 

79 returns None. 

80 """ 

81 

82 # Run the Python preprocessor to convert Python 2 to Python 3 

83 source_filename = Path(pyfile_path).name 

84 python2_source = Path(pyfile_path).read_text() 

85 python3_source = preprocess(python2_source) 

86 with TemporaryDirectory() as tmp: 

87 new_path = Path(tmp) / source_filename 

88 new_path.write_text(python3_source) 

89 pyfile_path = str(new_path) 

90 

91 if not tree_sitter: 

92 # Open Python file as a giant string 

93 with open(pyfile_path) as f: 

94 file_contents = f.read() 

95 

96 file_name = pyfile_path.split("/")[-1] 

97 

98 # Count the number of lines in the file 

99 with open(pyfile_path) as f: 

100 file_list = f.readlines() 

101 line_count = len(file_list) 

102 

103 # Create a PyASTToCAST Object 

104 if legacy: 

105 convert = py_ast_to_cast.PyASTToCAST(file_name, legacy=True) 

106 else: 

107 convert = py_ast_to_cast.PyASTToCAST(file_name) 

108 

109 # 'Root' the current working directory so that it's where the 

110 # Source file we're generating CAST for is (for Import statements) 

111 old_path = os.getcwd() 

112 try: 

113 idx = pyfile_path.rfind("/") 

114 

115 if idx > -1: 

116 curr_path = pyfile_path[0:idx] 

117 os.chdir(curr_path) 

118 else: 

119 curr_path = "./" + pyfile_path 

120 

121 # Parse the Python program's AST and create the CAST 

122 contents = ast.parse(file_contents) 

123 C = convert.visit(contents, {}, {}) 

124 C.source_refs = [SourceRef(file_name, None, None, 1, line_count)] 

125 finally: 

126 os.chdir(old_path) 

127 

128 out_cast = cast.CAST([C], "python") 

129 else: 

130 file_name = pyfile_path.split("/")[-1] 

131 out_cast = TS2CAST(pyfile_path).out_cast 

132 

133 if agraph: 

134 V = CASTToAGraphVisitor(out_cast) 

135 last_slash_idx = file_name.rfind("/") 

136 file_ending_idx = file_name.rfind(".") 

137 pdf_file_name = ( 

138 f"{file_name[last_slash_idx + 1 : file_ending_idx]}.pdf" 

139 ) 

140 V.to_pdf(pdf_file_name) 

141 

142 # Then, print CAST as JSON 

143 if cast_obj: 

144 return out_cast 

145 else: 

146 if rawjson: 

147 print( 

148 json.dumps( 

149 out_cast.to_json_object(), sort_keys=True, indent=None 

150 ) 

151 ) 

152 else: 

153 if std_out: 

154 print(out_cast.to_json_str()) 

155 else: 

156 out_name = file_name.split(".")[0] 

157 print("Writing CAST to " + out_name + "--CAST.json") 

158 out_handle = open(out_name + "--CAST.json", "w") 

159 out_handle.write(out_cast.to_json_str()) 

160 

161 

162if __name__ == "__main__": 

163 args = get_args() 

164 python_to_cast( 

165 args.pyfile_path, 

166 args.agraph, 

167 args.stdout, 

168 args.rawjson, 

169 args.legacy, 

170 tree_sitter=args.ts 

171 )