Coverage for src/prisma/cli/utils.py: 82%

1from __future__ import annotations 

2 

3import os 

4import sys 

5import logging 

6from enum import Enum 

7from typing import ( 

8 Any, 

9 List, 

10 Type, 

11 Union, 

12 Mapping, 

13 NoReturn, 

14 Optional, 

15 overload, 

16) 

17from pathlib import Path 

18from typing_extensions import override 

19 

20import click 

21 

22from . import prisma 

23from ..utils import module_exists 

24from .._types import Literal 

25 

26log: logging.Logger = logging.getLogger(__name__) 

27 

28 

29class PrismaCLI(click.MultiCommand): 

30 base_package: str = 'prisma.cli.commands' 

31 folder: Path = Path(__file__).parent / 'commands' 

32 

33 @override 

34 def list_commands(self, ctx: click.Context) -> List[str]: # noqa: ARG002 

35 commands: List[str] = [] 

36 

37 for path in self.folder.iterdir(): 

38 name = path.name 

39 if name.startswith('_'): 

40 continue 

41 

42 if name.endswith('.py'): 

43 commands.append(path.stem) 

44 elif is_module(path): 44 ↛ 37line 44 didn't jump to line 37, because the condition on line 44 was never false

45 commands.append(name) 

46 

47 commands.sort() 

48 return commands 

49 

50 @override 

51 def get_command(self, ctx: click.Context, cmd_name: str) -> Optional[click.Command]: # noqa: ARG002 

52 name = f'{self.base_package}.{cmd_name}' 

53 if not module_exists(name): 

54 # command not found 

55 return None 

56 

57 mod = __import__(name, None, None, ['cli']) 

58 

59 assert hasattr(mod, 'cli'), f'Expected command module {name} to contain a "cli" attribute' 

60 assert isinstance(mod.cli, click.Command), ( 

61 f'Expected command module attribute {name}.cli to be a {click.Command} ' 

62 f'instance but got {type(mod.cli)} instead' 

63 ) 

64 

65 return mod.cli 

66 

67 

68class PathlibPath(click.Path): 

69 """A Click path argument that returns a pathlib Path, not a string""" 

70 

71 @override 

72 def convert( 

73 self, 

74 value: str | os.PathLike[str], 

75 param: click.Parameter | None, 

76 ctx: click.Context | None, 

77 ) -> Path: 

78 return Path(str(super().convert(value, param, ctx))) 

79 

80 

81class EnumChoice(click.Choice): 

82 """A Click choice argument created from an Enum 

83 

84 choices are gathered from enum values, not their python keys, e.g. 

85 

86 class MyEnum(str, Enum): 

87 foo = 'bar' 

88 

89 results in click.Choice(['bar']) 

90 """ 

91 

92 def __init__(self, enum: Type[Enum]) -> None: 

93 if str not in enum.__mro__: 

94 raise TypeError('Enum does not subclass `str`') 

95 

96 self.__enum = enum 

97 super().__init__([item.value for item in enum.__members__.values()]) 

98 

99 @override 

100 def convert( 

101 self, 

102 value: str, 

103 param: Optional[click.Parameter], 

104 ctx: Optional[click.Context], 

105 ) -> str: 

106 return str(self.__enum(super().convert(value, param, ctx)).value) 

107 

108 

109def is_module(path: Path) -> bool: 

110 return path.is_dir() and path.joinpath('__init__.py').exists() 

111 

112 

113def maybe_exit(retcode: int) -> None: 

114 """Exit if given a non-zero exit code""" 

115 if retcode != 0: 115 ↛ 116line 115 didn't jump to line 116, because the condition on line 115 was never true

116 sys.exit(retcode) 

117 

118 

119def generate_client(schema: Optional[str] = None, *, reload: bool = False) -> None: 

120 """Run `prisma generate` and update sys.modules""" 

121 args = ['generate'] 

122 if schema is not None: 

123 args.append(f'--schema={schema}') 

124 

125 maybe_exit(prisma.run(args)) 

126 

127 if reload: 

128 for name in sys.modules.copy(): 

129 if 'prisma' in name and 'generator' not in name: 

130 sys.modules.pop(name, None) 

131 

132 

133def warning(message: str) -> None: 

134 click.echo(click.style('WARNING: ', fg='bright_yellow') + click.style(message, bold=True)) 

135 

136 

137@overload 

138def error(message: str) -> NoReturn: ... 

139 

140 

141@overload 

142def error(message: str, exit_: Literal[True]) -> NoReturn: ... 

143 

144 

145@overload 

146def error(message: str, exit_: Literal[False]) -> None: ... 

147 

148 

149def error(message: str, exit_: bool = True) -> Union[None, NoReturn]: 

150 click.echo(click.style(message, fg='bright_red', bold=True), err=True) 

151 if exit_: 

152 sys.exit(1) 

153 else: 

154 return None 

155 

156 

157def pretty_info(mapping: Mapping[str, Any]) -> str: 

158 """Pretty print a mapping 

159 

160 e.g {'foo': 'bar', 'hello': 1} 

161 

162 foo : bar 

163 hello : 1 

164 """ 

165 pad = max(len(k) for k in mapping.keys()) 

166 return '\n'.join(f'{k.ljust(pad)} : {v}' for k, v in mapping.items())