Coverage for src / competitive_verifier / models / shell.py: 100%

1import pathlib 

2from subprocess import CompletedProcess 

3from typing import Annotated, Literal, overload 

4 

5from pydantic import BaseModel, Field 

6from typing_extensions import TypeAliasType 

7 

8from competitive_verifier.exec import exec_command 

9 

10 

11class ShellCommand(BaseModel): 

12 command: list[str] | str = Field( 

13 description="Shell command", 

14 ) 

15 """Shell command 

16 """ 

17 

18 env: dict[str, str] | None = Field( 

19 default=None, 

20 description="Envitonment variables for command", 

21 ) 

22 """Envitonment variables for command 

23 """ 

24 

25 cwd: pathlib.Path | None = Field( 

26 default=None, 

27 description="The working directory of child process.", 

28 ) 

29 """The working directory of child process. 

30 """ 

31 

32 @overload 

33 def exec_command( 

34 self, 

35 *, 

36 text: Literal[False] = False, 

37 check: bool = False, 

38 capture_output: bool = False, 

39 group_log: bool = False, 

40 ) -> CompletedProcess[bytes]: ... 

41 

42 @overload 

43 def exec_command( 

44 self, 

45 *, 

46 text: Literal[True], 

47 check: bool = False, 

48 capture_output: bool = False, 

49 group_log: bool = False, 

50 ) -> CompletedProcess[str]: ... 

51 

52 def exec_command( 

53 self, 

54 *, 

55 text: bool = False, 

56 check: bool = False, 

57 capture_output: bool = False, 

58 group_log: bool = False, 

59 ) -> CompletedProcess[str] | CompletedProcess[bytes]: 

60 return exec_command( 

61 command=self.command, 

62 env=self.env, 

63 text=text, 

64 check=check, 

65 capture_output=capture_output, 

66 cwd=self.cwd, 

67 group_log=group_log, 

68 ) 

69 

70 @classmethod 

71 def parse_command_like(cls, cmd: "ShellCommandLike") -> "ShellCommand": 

72 if isinstance(cmd, (str, list)): 

73 return ShellCommand(command=cmd) 

74 return cmd 

75 

76 

77ShellCommandLike = TypeAliasType( 

78 "ShellCommandLike", 

79 Annotated[ 

80 ShellCommand | list[str] | str, 

81 Field( 

82 examples=[ 

83 "command", 

84 ["command", "arg1", "arg2"], 

85 ShellCommand( 

86 command=["command", "arg1", "arg2"], 

87 env={"ENVVAR": "DUMMY"}, 

88 cwd=pathlib.Path("/work"), 

89 ), 

90 ShellCommand( 

91 command="command", 

92 env={"ENVVAR": "DUMMY"}, 

93 cwd=pathlib.Path("/work"), 

94 ), 

95 ] 

96 ), 

97 ], 

98)