1 import ast

2 import inspect

3 import textwrap

4 import tokenize

5 import types

6 import warnings

7 from bisect import bisect_right

8 from typing import Iterable

9 from typing import Iterator

10 from typing import List

11 from typing import Optional

12 from typing import overload

13 from typing import Tuple

14 from typing import Union

15

16

17 class Source:

18 """An immutable object holding a source code fragment.

19

20 When using Source(...), the source lines are deindented.

21 """

22

23 def __init__(self, obj: object = None) -> None:

24 if not obj:

25 self.lines: List[str] = []

26 elif isinstance(obj, Source):

27 self.lines = obj.lines

28 elif isinstance(obj, (tuple, list)):

29 self.lines = deindent(x.rstrip("\n") for x in obj)

30 elif isinstance(obj, str):

31 self.lines = deindent(obj.split("\n"))

32 else:

33 try:

34 rawcode = getrawcode(obj)

35 src = inspect.getsource(rawcode)

36 except TypeError:

37 src = inspect.getsource(obj) # type: ignore[arg-type]

38 self.lines = deindent(src.split("\n"))

39

40 def __eq__(self, other: object) -> bool:

41 if not isinstance(other, Source):

42 return NotImplemented

43 return self.lines == other.lines

44

45 # Ignore type because of https://github.com/python/mypy/issues/4266.

46 __hash__ = None # type: ignore

47

48 @overload

49 def __getitem__(self, key: int) -> str:

50 ...

51

52 @overload

53 def __getitem__(self, key: slice) -> "Source":

54 ...

55

56 def __getitem__(self, key: Union[int, slice]) -> Union[str, "Source"]:

57 if isinstance(key, int):

58 return self.lines[key]

59 else:

60 if key.step not in (None, 1):

61 raise IndexError("cannot slice a Source with a step")

62 newsource = Source()

63 newsource.lines = self.lines[key.start : key.stop]

64 return newsource

65

66 def __iter__(self) -> Iterator[str]:

67 return iter(self.lines)

68

69 def __len__(self) -> int:

70 return len(self.lines)

71

72 def strip(self) -> "Source":

73 """Return new Source object with trailing and leading blank lines removed."""

74 start, end = 0, len(self)

75 while start < end and not self.lines[start].strip():

76 start += 1

77 while end > start and not self.lines[end - 1].strip():

78 end -= 1

79 source = Source()

80 source.lines[:] = self.lines[start:end]

81 return source

82

83 def indent(self, indent: str = " " * 4) -> "Source":

84 """Return a copy of the source object with all lines indented by the

85 given indent-string."""

86 newsource = Source()

87 newsource.lines = [(indent + line) for line in self.lines]

88 return newsource

89

90 def getstatement(self, lineno: int) -> "Source":

91 """Return Source statement which contains the given linenumber

92 (counted from 0)."""

93 start, end = self.getstatementrange(lineno)

94 return self[start:end]

95

96 def getstatementrange(self, lineno: int) -> Tuple[int, int]:

97 """Return (start, end) tuple which spans the minimal statement region

98 which containing the given lineno."""

99 if not (0 <= lineno < len(self)):

100 raise IndexError("lineno out of range")

101 ast, start, end = getstatementrange_ast(lineno, self)

102 return start, end

103

104 def deindent(self) -> "Source":

105 """Return a new Source object deindented."""

106 newsource = Source()

107 newsource.lines[:] = deindent(self.lines)

108 return newsource

109

110 def __str__(self) -> str:

111 return "\n".join(self.lines)

112

113

114#

115# helper functions

116#

117

118

119 def findsource(obj) -> Tuple[Optional[Source], int]:

120 try:

121 sourcelines, lineno = inspect.findsource(obj)

122 except Exception:

123 return None, -1

124 source = Source()

125 source.lines = [line.rstrip() for line in sourcelines]

126 return source, lineno

127

128

129 def getrawcode(obj: object, trycall: bool = True) -> types.CodeType:

130 """Return code object for given function."""

131 try:

132 return obj.__code__ # type: ignore[attr-defined,no-any-return]

133 except AttributeError:

134 pass

135 if trycall:

136 call = getattr(obj, "__call__", None)

137 if call and not isinstance(obj, type):

138 return getrawcode(call, trycall=False)

139 raise TypeError(f"could not get code object for {obj!r}")

140

141

142 def deindent(lines: Iterable[str]) -> List[str]:

143 return textwrap.dedent("\n".join(lines)).splitlines()

144

145

146 def get_statement_startend2(lineno: int, node: ast.AST) -> Tuple[int, Optional[int]]:

147 # Flatten all statements and except handlers into one lineno-list.

148 # AST's line numbers start indexing at 1.

149 values: List[int] = []

150 for x in ast.walk(node):

151 if isinstance(x, (ast.stmt, ast.ExceptHandler)):

152 # Before Python 3.8, the lineno of a decorated class or function pointed at the decorator.

153 # Since Python 3.8, the lineno points to the class/def, so need to include the decorators.

154 if isinstance(x, (ast.ClassDef, ast.FunctionDef, ast.AsyncFunctionDef)):

155 for d in x.decorator_list:

156 values.append(d.lineno - 1)

157 values.append(x.lineno - 1)

158 for name in ("finalbody", "orelse"):

159 val: Optional[List[ast.stmt]] = getattr(x, name, None)

160 if val:

161 # Treat the finally/orelse part as its own statement.

162 values.append(val[0].lineno - 1 - 1)

163 values.sort()

164 insert_index = bisect_right(values, lineno)

165 start = values[insert_index - 1]

166 if insert_index >= len(values):

167 end = None

168 else:

169 end = values[insert_index]

170 return start, end

171

172

173 def getstatementrange_ast(

174 lineno: int,

175 source: Source,

176 assertion: bool = False,

177 astnode: Optional[ast.AST] = None,

178) -> Tuple[ast.AST, int, int]:

179 if astnode is None:

180 content = str(source)

181 # See #4260:

182 # Don't produce duplicate warnings when compiling source to find AST.

183 with warnings.catch_warnings():

184 warnings.simplefilter("ignore")

185 astnode = ast.parse(content, "source", "exec")

186

187 start, end = get_statement_startend2(lineno, astnode)

188 # We need to correct the end:

189 # - ast-parsing strips comments

190 # - there might be empty lines

191 # - we might have lesser indented code blocks at the end

192 if end is None:

193 end = len(source.lines)

194

195 if end > start + 1:

196 # Make sure we don't span differently indented code blocks

197 # by using the BlockFinder helper used which inspect.getsource() uses itself.

198 block_finder = inspect.BlockFinder()

199 # If we start with an indented line, put blockfinder to "started" mode.

200 block_finder.started = source.lines[start][0].isspace()

201 it = ((x + "\n") for x in source.lines[start:end])

202 try:

203 for tok in tokenize.generate_tokens(lambda: next(it)):

204 block_finder.tokeneater(*tok)

205 except (inspect.EndOfBlock, IndentationError):

206 end = block_finder.last + start

207 except Exception:

208 pass

209

210 # The end might still point to a comment or empty line, correct it.

211 while end:

212 line = source.lines[end - 1].lstrip()

213 if line.startswith("#") or not line:

214 end -= 1

215 else:

216 break

217 return astnode, start, end