Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / codegen / astkit / edits.py: 31%

1# Copyright 2026 Huawei Technologies Co., Ltd 

2# Licensed under the Apache License, Version 2.0 

3# ============================================================================ 

4"""Positional text edits over a modeling source. 

5 

6The forward lowerer (emit/parallel) turns a collection of structural plan 

7(index.py) into a list of ``TextEdit``s and replays them. All offsets are byte 

8offsets into the *original* text; ``apply_edits`` validates ordering and never 

9edits the same span twice, so a malformed edit set fails loudly instead of 

10silently corrupting the file. 

11 

12Only the edits are positional; every helper that builds an edit from a 

13``FunctionInfo`` / span does the spelling itself. Nothing here imports torch. 

14""" 

15from __future__ import annotations 

16 

17from dataclasses import dataclass 

18from typing import Optional, Sequence 

19 

20 

21@dataclass 

22class TextEdit: 

23 """One ``[start, end)`` byte-span replacement in a source string.""" 

24 

25 start: int 

26 end: int 

27 replacement: str 

28 

29 def __post_init__(self) -> None: 

30 if self.start < 0 or self.end < self.start: 

31 raise ValueError( 

32 f"TextEdit has an invalid span: start={self.start} end={self.end}" 

33 ) 

34 

35 

36def apply_edits(text: str, edits: Sequence[TextEdit]) -> str: 

37 """Replay ``edits`` on ``text``, earliest-first for stable ordering. 

38 

39 Edits are authored against the *original* text. Applying one splice shifts 

40 every later offset by the edit's length delta, so each edit is re-anchored 

41 by the cumulative drift of the edits before it. Edits must be 

42 non-overlapping and stay within the original text — checked up front. 

43 """ 

44 ordered = sorted(edits, key=lambda e: (e.start, e.end)) 

45 if ordered: 

46 prev_end = ordered[0].start 

47 for edit in ordered: 

48 if edit.start < prev_end: 

49 raise ValueError( 

50 f"apply_edits: overlapping edits — [{edit.start}, {edit.end}) " 

51 f"overlaps the previous edit ending at {prev_end}" 

52 ) 

53 if edit.end > len(text): 

54 raise ValueError( 

55 f"apply_edits: edit [{edit.start}, {edit.end}) exceeds text " 

56 f"length {len(text)}" 

57 ) 

58 prev_end = edit.end 

59 

60 out = list(text) 

61 shift = 0 

62 for edit in ordered: 

63 start = edit.start + shift 

64 end = edit.end + shift 

65 replacement = edit.replacement 

66 out[start:end] = replacement 

67 shift += len(replacement) - (edit.end - edit.start) 

68 return "".join(out) 

69 

70 

71def replace_function_body( 

72 source_text: str, 

73 func: "FunctionInfo", 

74 new_body: str, 

75 *, 

76 indent: str = " ", 

77) -> tuple[TextEdit, ...]: 

78 """Build the edit that swaps ``func``'s body for ``new_body``. 

79 

80 ``new_body`` is the replacement body WITHOUT line indentation — every line 

81 is re-indented to the method's existing indent (``func``'s first body 

82 statement). A docstring inside the body is preserved by the caller (the 

83 lowerer emits it explicitly). A single-line/pass body has no replaceable 

84 span: the whole method is rewritten as ``def <name>(...):\n<indented body>`` 

85 from the file text's own signature line. 

86 

87 Returns an empty tuple when there is nothing to replace (a function with 

88 no body span and no way to anchor). 

89 """ 

90 if func.body_start is None or func.body_end is None: 

91 return () 

92 # Re-indent per line, preserving blank lines and inner indentation. The 

93 # caller renders ``new_body`` at column-0-relative depth (blank lines carry 

94 # no indentation), so a statement's relative nesting survives; only the 

95 # method's own ``indent`` is applied once. 

96 body_lines = [ 

97 (indent + line) if line.strip() else "" 

98 for line in new_body.splitlines(keepends=True) 

99 ] 

100 replacement = "".join(body_lines) 

101 return (TextEdit(func.body_start, func.body_end, replacement),) 

102 

103 

104def insert_after_imports( 

105 source_text: str, 

106 import_end: Optional[int], 

107 insert_text: str, 

108 *, 

109 default_offset: int = 0, 

110) -> Optional[TextEdit]: 

111 """Build an edit inserting ``insert_text`` right after the import block. 

112 

113 ``import_end`` is the byte offset just past the last import statement. A 

114 ``None`` ``import_end`` means no import block: fall back to 

115 ``default_offset`` (the docstring end or 0). Returns ``None`` when there is 

116 no sensible anchor — caller treats that as "insert at the top". 

117 """ 

118 anchor = import_end if import_end is not None else default_offset 

119 if anchor < 0 or anchor > len(source_text): 

120 return None 

121 # Keep the inserted text on its own lines, bracketed by newlines, so it 

122 # never glues onto a trailing comment or the first statement. 

123 return TextEdit(anchor, anchor, "\n" + insert_text + "\n") 

124 

125 

126__all__ = [ 

127 "TextEdit", 

128 "apply_edits", 

129 "insert_after_imports", 

130 "replace_function_body", 

131]