Coverage for src/workstack/core/file_utils.py: 80%

1"""File operation utilities.""" 

2 

3import os 

4import shutil 

5import tempfile 

6from collections.abc import Iterator 

7from contextlib import contextmanager 

8from pathlib import Path 

9 

10 

11def _is_critical_system_file(path: Path) -> bool: 

12 """Check if path is a critical system file that should not be modified by tests. 

13 

14 Returns True for shell rc files in the user's real home directory. 

15 This is a safety check to prevent tests from accidentally modifying real system files. 

16 """ 

17 if not path.exists(): 

18 return False 

19 

20 real_home = Path.home() 

21 critical_files = { 

22 real_home / ".bashrc", 

23 real_home / ".zshrc", 

24 real_home / ".bash_profile", 

25 real_home / ".profile", 

26 real_home / ".config" / "fish" / "config.fish", 

27 } 

28 

29 try: 

30 resolved_path = path.resolve() 

31 return resolved_path in critical_files 

32 except (OSError, RuntimeError): 

33 return False 

34 

35 

36@contextmanager 

37def atomic_write(target_path: Path, *, mode: str = "w", encoding: str = "utf-8") -> Iterator: 

38 """Write to a file atomically using a temporary file and rename. 

39 

40 This context manager ensures that file writes are atomic - either the entire 

41 write succeeds or the original file is left untouched. It also preserves 

42 file permissions from the original file if it exists. 

43 

44 Args: 

45 target_path: Final destination path for the file 

46 mode: File open mode (default "w") 

47 encoding: File encoding (default "utf-8") 

48 

49 Yields: 

50 File handle for writing 

51 

52 Example: 

53 with atomic_write(Path("config.txt")) as f: 

54 f.write("new content") 

55 

56 Note: Exception handling for cleanup is acceptable here per EXCEPTION_HANDLING.md 

57 as this is encapsulating necessary exception handling at an error boundary. 

58 

59 Raises: 

60 RuntimeError: If attempting to write to a critical system file (safety check) 

61 """ 

62 target_path = Path(target_path) 

63 

64 # Safety check: prevent accidental writes to critical system files 

65 if _is_critical_system_file(target_path): 

66 raise RuntimeError( 

67 f"Refusing to write to critical system file: {target_path}. " 

68 "This is a safety check to prevent accidental modification of shell rc files. " 

69 "If you're writing a test, use a temporary directory instead of Path.home()." 

70 ) 

71 

72 target_path.parent.mkdir(parents=True, exist_ok=True) 

73 

74 # Create temp file in same directory to ensure same filesystem 

75 temp_fd, temp_path = tempfile.mkstemp( 

76 dir=target_path.parent, prefix=f".{target_path.name}.", suffix=".tmp" 

77 ) 

78 

79 try: 

80 with os.fdopen(temp_fd, mode, encoding=encoding) as f: 

81 yield f 

82 

83 # Preserve permissions from original file if it exists 

84 if target_path.exists(): 

85 shutil.copystat(target_path, temp_path) 

86 

87 # Atomic rename 

88 os.rename(temp_path, target_path) 

89 

90 except OSError: 

91 # Clean up temp file on error - acceptable use of exception handling 

92 # per EXCEPTION_HANDLING.md (cleanup during error boundaries) 

93 try: 

94 os.unlink(temp_path) 

95 except FileNotFoundError: 

96 pass # File was never created if mkstemp failed 

97 raise 

98 

99 

100def extract_plan_title(plan_path: Path) -> str | None: 

101 """Extract the first heading from a markdown plan file. 

102 

103 Uses python-frontmatter library to properly parse YAML frontmatter, 

104 then extracts the first line starting with # from the content. 

105 

106 Args: 

107 plan_path: Path to the .PLAN.md file 

108 

109 Returns: 

110 The heading text (without the # prefix), or None if not found or file doesn't exist 

111 """ 

112 if not plan_path.exists(): 

113 return None 

114 

115 import frontmatter 

116 

117 # Parse file with frontmatter library (handles YAML frontmatter properly) 

118 post = frontmatter.load(str(plan_path)) 

119 

120 # Get the content (without frontmatter) 

121 content = post.content 

122 lines = content.splitlines() 

123 

124 # Find first heading 

125 for line in lines: 

126 stripped = line.strip() 

127 if stripped.startswith("#"): 

128 # Remove all # symbols and strip whitespace 

129 title = stripped.lstrip("#").strip() 

130 if title: 

131 return title 

132 

133 return None