Coverage for src/pydal2sql/cli.py: 100%

28 statements  

« prev     ^ index     » next       coverage.py v7.2.7, created at 2023-07-20 20:24 +0200

1""" 

2CLI tool to generate SQL from PyDAL code. 

3""" 

4 

5import argparse 

6import pathlib 

7import select 

8import string 

9import sys 

10import textwrap 

11from typing import IO, Optional 

12 

13import rich 

14from rich.prompt import Prompt 

15 

16from .helpers import flatten 

17from .magic import find_missing_variables, generate_magic_code 

18 

19 

20class PrettyParser(argparse.ArgumentParser): # pragma: no cover 

21 """ 

22 Add 'rich' to the argparse output. 

23 """ 

24 

25 def _print_message(self, message: str, file: Optional[IO[str]] = None) -> None: 

26 rich.print(message, file=file) 

27 

28 

29def has_stdin_data() -> bool: # pragma: no cover 

30 """ 

31 Check if the program starts with cli data (pipe | or redirect ><). 

32 

33 See Also: 

34 https://stackoverflow.com/questions/3762881/how-do-i-check-if-stdin-has-some-data 

35 """ 

36 return any( 

37 select.select( 

38 [ 

39 sys.stdin, 

40 ], 

41 [], 

42 [], 

43 0.0, 

44 )[0] 

45 ) 

46 

47 

48def handle_cli( 

49 code: str, 

50 db_type: str = None, 

51 tables: list[list[str]] = None, 

52 verbose: bool = False, 

53 noop: bool = False, 

54 magic: bool = False, 

55) -> None: 

56 """ 

57 Handle user input. 

58 """ 

59 to_execute = string.Template( 

60 textwrap.dedent( 

61 """ 

62 from pydal import * 

63 from pydal.objects import * 

64 from pydal.validators import * 

65 

66 from pydal2sql import generate_sql 

67 

68 db = database = DAL(None, migrate=False) 

69 

70 tables = $tables 

71 db_type = '$db_type' 

72 

73 $extra 

74 

75 $code 

76 

77 if not tables: 

78 tables = db._tables 

79 

80 for table in tables: 

81 print(generate_sql(db[table], db_type)) 

82 """ 

83 ) 

84 ) 

85 

86 generated_code = to_execute.substitute( 

87 {"tables": flatten(tables or []), "db_type": db_type or "", "code": textwrap.dedent(code), "extra": ""} 

88 ) 

89 if verbose or noop: 

90 rich.print(generated_code, file=sys.stderr) 

91 

92 if not noop: 

93 try: 

94 exec(generated_code) # nosec: B102 

95 except NameError: 

96 # something is missing! 

97 missing_vars = find_missing_variables(generated_code) 

98 if not magic: 

99 rich.print( 

100 f"Your code is missing some variables: {missing_vars}. Add these or try --magic", file=sys.stderr 

101 ) 

102 else: 

103 extra_code = generate_magic_code(missing_vars) 

104 

105 generated_code = to_execute.substitute( 

106 { 

107 "tables": flatten(tables or []), 

108 "db_type": db_type or "", 

109 "extra": extra_code, 

110 "code": textwrap.dedent(code), 

111 } 

112 ) 

113 

114 if verbose: 

115 print(generated_code, file=sys.stderr) 

116 

117 exec(generated_code) # nosec: B102 

118 

119 

120def app() -> None: # pragma: no cover 

121 """ 

122 Entrypoint for the pydal2sql cli command. 

123 """ 

124 parser = PrettyParser( 

125 prog="pydal2sql", 

126 formatter_class=argparse.RawDescriptionHelpFormatter, 

127 description="""[green]CLI tool to generate SQL from PyDAL code.[/green]\n 

128 Aside from using cli arguments, you can also configure the tool in your code. 

129 You can set the following variables: 

130 

131 db_type: str = 'sqlite' # your desired database type; 

132 tables: list[str] = [] # your desired tables to generate SQL for;""", 

133 epilog="Example: [i]cat models.py | pydal2sql sqlite[/i]", 

134 ) 

135 

136 parser.add_argument("filename", nargs="?", help="Which file to load? Can also be done with stdin.") 

137 

138 parser.add_argument( 

139 "db_type", nargs="?", help="Which database dialect to generate ([blue]postgres, sqlite, mysql[/blue])" 

140 ) 

141 

142 parser.add_argument("--verbose", "-v", help="Show more info", action=argparse.BooleanOptionalAction, default=False) 

143 

144 parser.add_argument( 

145 "--noop", "-n", help="Only show code, don't run it.", action=argparse.BooleanOptionalAction, default=False 

146 ) 

147 

148 parser.add_argument( 

149 "--magic", "-m", help="Perform magic to fix missing vars.", action=argparse.BooleanOptionalAction, default=False 

150 ) 

151 

152 parser.add_argument( 

153 "-t", 

154 "--table", 

155 "--tables", 

156 action="append", 

157 nargs="+", 

158 help="One or more tables to generate. By default, all tables in the file will be used.", 

159 ) 

160 

161 args = parser.parse_args() 

162 

163 db_type = args.db_type or args.filename 

164 

165 load_file_mode = (filename := args.filename) and filename.endswith(".py") 

166 

167 if not (has_stdin_data() or load_file_mode): 

168 if not db_type: 

169 db_type = Prompt.ask("Which database type do you want to use?", choices=["sqlite", "postgres", "mysql"]) 

170 

171 rich.print("Please paste your define tables code below and press ctrl-D when finished.", file=sys.stderr) 

172 

173 # else: data from stdin 

174 # py code or cli args should define settings. 

175 if load_file_mode: 

176 db_type = args.db_type 

177 text = pathlib.Path(filename).read_text() 

178 else: 

179 text = sys.stdin.read() 

180 

181 return handle_cli(text, db_type, args.table, verbose=args.verbose, noop=args.noop, magic=args.magic)