Coverage for src/duty/context.py: 94.64%

1"""Module containing the context definition.""" 

2 

3from __future__ import annotations 

4 

5import os 

6from contextlib import contextmanager, suppress 

7from typing import TYPE_CHECKING, Any, Callable, Union 

8 

9from failprint.runners import run as failprint_run 

10 

11from duty.exceptions import DutyFailure 

12from duty.tools import Tool 

13 

14if TYPE_CHECKING: 

15 from collections.abc import Iterator 

16 

17CmdType = Union[str, list[str], Callable] 

18 

19 

20class Context: 

21 """A simple context class. 

22 

23 Context instances are passed to functions decorated with `duty`. 

24 """ 

25 

26 def __init__(self, options: dict[str, Any], options_override: dict[str, Any] | None = None) -> None: 

27 """Initialize the context. 

28 

29 Parameters: 

30 options: Base options specified in `@duty(**options)`. 

31 options_override: Options that override `run` and `@duty` options. 

32 This argument is used to allow users to override options from the CLI or environment. 

33 """ 

34 self._options = options 

35 self._option_stack: list[dict[str, Any]] = [] 

36 self._options_override = options_override or {} 

37 

38 @contextmanager 

39 def cd(self, directory: str) -> Iterator: 

40 """Change working directory as a context manager. 

41 

42 Parameters: 

43 directory: The directory to go into. 

44 

45 Yields: 

46 Nothing. 

47 """ 

48 if not directory: 

49 yield 

50 return 

51 old_wd = os.getcwd() 

52 os.chdir(directory) 

53 try: 

54 yield 

55 finally: 

56 os.chdir(old_wd) 

57 

58 def run(self, cmd: CmdType, **options: Any) -> str: 

59 """Run a command in a subprocess or a Python callable. 

60 

61 Parameters: 

62 cmd: A command or a Python callable. 

63 options: Options passed to `failprint` functions. 

64 

65 Raises: 

66 DutyFailure: When the exit code / function result is greather than 0. 

67 

68 Returns: 

69 The output of the command. 

70 """ 

71 final_options = dict(self._options) 

72 final_options.update(options) 

73 

74 if "command" not in final_options and isinstance(cmd, Tool): 74 ↛ 75line 74 didn't jump to line 75 because the condition on line 74 was never true

75 with suppress(ValueError): 

76 final_options["command"] = cmd.cli_command 

77 

78 allow_overrides = final_options.pop("allow_overrides", True) 

79 workdir = final_options.pop("workdir", None) 

80 

81 if allow_overrides: 

82 final_options.update(self._options_override) 

83 

84 with self.cd(workdir): 

85 try: 

86 result = failprint_run(cmd, **final_options) 

87 except KeyboardInterrupt as ki: 

88 raise DutyFailure(130) from ki 

89 

90 if result.code: 

91 raise DutyFailure(result.code) 

92 

93 return result.output 

94 

95 @contextmanager 

96 def options(self, **opts: Any) -> Iterator: 

97 """Change options as a context manager. 

98 

99 Can be nested as will, previous options will pop once out of the with clause. 

100 

101 Parameters: 

102 **opts: Options used in `run`. 

103 

104 Yields: 

105 Nothing. 

106 """ 

107 self._option_stack.append(self._options) 

108 self._options = {**self._options, **opts} 

109 try: 

110 yield 

111 finally: 

112 self._options = self._option_stack.pop()