Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / codegen / astkit / index.py: 26%
140 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd
2# Licensed under the Apache License, Version 2.0
3# ============================================================================
4"""Index a modeling source: locate classes, forwards, and import regions.
6``lower_forward_boundaries`` (emit/parallel) rewrites a boundary class's
7``forward`` in place. To do that it needs three answers from the source:
91. where a top-level ``class`` starts (and how far it spans) so an edit never
10 bleeds into the next class;
112. where that class's ``forward`` sits — its indentation, its first statement,
12 and the byte span of its body (the part that gets replaced);
133. where the module's import block ends, so the lowerer can insert the frozen
14 per-class boundary constants *before* the first class body but after the
15 imports + module docstring (PEP 8: imports first).
17Everything is computed with ``ast`` and recorded as 1-based line numbers +
18character offsets into the *original* ``source_text``. Editing happens on the
19unmodified text (the index is never reused against a patched string), so the
20offset mapping stays valid for the whole call.
22Nothing here imports torch or the training stack.
23"""
24from __future__ import annotations
26import ast
27import textwrap
28from dataclasses import dataclass, field
29from typing import Optional
32@dataclass
33class FunctionInfo:
34 """Location of one method inside a class body."""
36 name: str
37 #: Byte offset of the first statement in the function body. ``None`` if the
38 #: body is just ``pass`` / ``...`` on the same line as ``def`` (single-line
39 #: function) — there is nothing to replace.
40 body_start: Optional[int]
41 #: Byte offset just past the last statement of the body.
42 body_end: Optional[int]
43 #: Character offset of the ``def`` keyword (== offset of the first decorator
44 #: when one is present); used to delete a decorator along with the body.
45 def_offset: int
46 #: Positional parameter names of the signature (BOUND to the plan's
47 #: ``arg_name`` by ``sharding_applier._bind_input_indices``). ``self`` is
48 #: excluded; ``*args`` / ``**kwargs`` are not part of the positional list.
49 param_names: list[str]
50 #: Keyword-only parameter names (``def f(self, x, *, y=1)`` -> ``["y"]``).
51 kwonly_names: list[str] = field(default_factory=list)
52 #: Name of the ``*args`` parameter, if present.
53 vararg_name: Optional[str] = None
54 #: Name of the ``**kwargs`` parameter, if present.
55 kwarg_name: Optional[str] = None
58@dataclass
59class ClassInfo:
60 """Location of one top-level class definition."""
62 name: str
63 #: 1-based line of the ``class`` keyword (decorators excluded).
64 class_line: int
65 #: Byte offset of the ``class`` keyword (first decorator, if any).
66 class_offset: int
67 #: Byte offset just past the class body's last statement.
68 end_offset: int
69 #: Methods keyed by name (only top-level ``def`` in the class block).
70 methods: dict[str, FunctionInfo]
73@dataclass
74class SourceIndex:
75 """Structural map of one modeling source text."""
77 source_text: str
78 classes: dict[str, ClassInfo]
79 #: Byte offset where the import block ends and code/classes may begin.
80 import_end: Optional[int]
81 #: Byte offset of the module docstring's closing quote (``None`` if absent).
82 docstring_end: Optional[int]
84 def find_class(self, name: str) -> Optional[ClassInfo]:
85 """Return the ``ClassInfo`` for a top-level class, or ``None``."""
86 return self.classes.get(name)
88 def find_forward(self, class_name: str) -> Optional[FunctionInfo]:
89 """Return the ``FunctionInfo`` for ``class_name.forward``, or ``None``.
91 A class without a ``forward`` (not a module) returns ``None``; the
92 caller decides whether that is an error.
93 """
94 info = self.classes.get(class_name)
95 if info is None:
96 return None
97 return info.methods.get("forward")
100def build_source_index(source_text: str) -> SourceIndex:
101 """Parse ``source_text`` and answer the three structural questions.
103 The index is built against the *original* text. Caller must apply edits to
104 the same unmodified string (edits.py ``apply_edits`` splices by offset).
105 """
106 tree = ast.parse(source_text)
108 line_offsets = _line_offsets(source_text)
110 def offset_for(line_no: int) -> int:
111 """Byte offset of the start of ``line_no`` (1-based)."""
112 if 1 <= line_no <= len(line_offsets):
113 return line_offsets[line_no - 1]
114 return len(source_text)
116 classes: dict[str, ClassInfo] = {}
117 import_end: Optional[int] = None
118 docstring_end: Optional[int] = None
120 body = tree.body
121 # Module docstring: the first statement is a constant string (PEP 257). Its
122 # closing quote is the first insertion point for appended constants, so
123 # Frozen constants belong after the docstring and before imports.
124 if body and isinstance(body[0], ast.Expr) and isinstance(body[0].value, ast.Constant):
125 if isinstance(body[0].value.value, str) and body[0].end_lineno is not None:
126 docstring_end = offset_for(body[0].end_lineno)
128 # Import block: the leading (or maximal) run of Import/ImportFrom at module
129 # level. In a well-formed module imports are first, so we take the span
130 # from the first import to the last import before the first non-import
131 # statement. Anything after the docstring but before the first import (a
132 # comment, ``__future__`` in a string, a module-level flag) is skipped.
133 for node in body:
134 if isinstance(node, (ast.Import, ast.ImportFrom)):
135 if node.end_lineno is not None:
136 import_end = offset_for(node.end_lineno)
138 for node in body:
139 if isinstance(node, ast.ClassDef):
140 class_info = _build_class_info(node, source_text, offset_for)
141 classes[class_info.name] = class_info
143 return SourceIndex(
144 source_text=source_text,
145 classes=classes,
146 import_end=import_end,
147 docstring_end=docstring_end,
148 )
151def _build_class_info(node: ast.ClassDef, text: str, offset_for) -> ClassInfo:
152 methods: dict[str, FunctionInfo] = {}
153 for child in node.body:
154 if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef)):
155 methods[child.name] = _build_function_info(child, text, offset_for)
156 end_offset = (
157 offset_for(node.end_lineno) if node.end_lineno is not None else len(text)
158 )
159 return ClassInfo(
160 name=node.name,
161 class_line=node.lineno,
162 class_offset=offset_for(node.lineno),
163 end_offset=end_offset,
164 methods=methods,
165 )
168def _build_function_info(node: ast.AST, text: str, offset_for) -> FunctionInfo:
169 """Collect offsets, parameters and ``*args``/``**kwargs`` names for a def node."""
170 decorators = getattr(node, "decorator_list", [])
171 first_def_line = node.lineno
172 if decorators:
173 first_def_line = min(decorators[0].lineno, first_def_line)
174 def_offset = offset_for(first_def_line)
176 # Positional parameters (``self`` excluded) — the names the boundary's
177 # ``arg_name`` binds against. Only POSITIONAL_ONLY / POSITIONAL_OR_KEYWORD
178 # are part of a positional call; ``*args`` / ``**kwargs`` are not.
179 params = getattr(node, "args", None)
180 positional: list[str] = []
181 kwonly: list[str] = []
182 vararg_name: Optional[str] = None
183 kwarg_name: Optional[str] = None
184 if params is not None:
185 for arg in params.posonlyargs + params.args:
186 if arg.arg != "self":
187 positional.append(arg.arg)
188 kwonly = [arg.arg for arg in params.kwonlyargs]
189 vararg_name = params.vararg.arg if params.vararg is not None else None
190 kwarg_name = params.kwarg.arg if params.kwarg is not None else None
192 body = node.body
193 body_start: Optional[int] = None
194 body_end: Optional[int] = None
195 # A function whose body is a single `pass` / `...` / single constant
196 # statement on the same line has no replaceable body span.
197 has_body = any(
198 not (
199 isinstance(stmt, ast.Pass)
200 or (isinstance(stmt, ast.Expr) and isinstance(stmt.value, ast.Constant) and stmt.value.value is Ellipsis)
201 )
202 for stmt in body
203 )
204 if has_body:
205 first_stmt = body[0]
206 body_start = offset_for(first_stmt.lineno)
207 last_stmt = body[-1]
208 # End just past the last statement's final character AND its trailing
209 # newline (when present). ``offset_for(end_lineno)`` is only the START
210 # of that line — using it here would leave the last statement (e.g.
211 # ``return x``) outside the replaceable span, so a body swap would
212 # splice a stray leftover statement after the replacement. The
213 # start-of-line offset plus the statement's column + 1 lands just past
214 # the line's newline; when the statement is the last line and the file
215 # has no trailing newline, that would overshoot EOF, so clamp to
216 # ``len(text)``.
217 body_end = (
218 min(offset_for(last_stmt.end_lineno) + last_stmt.end_col_offset + 1, len(text))
219 if last_stmt.end_lineno is not None
220 else len(text)
221 )
222 return FunctionInfo(
223 name=node.name,
224 body_start=body_start,
225 body_end=body_end,
226 def_offset=def_offset,
227 param_names=positional,
228 kwonly_names=kwonly,
229 vararg_name=vararg_name,
230 kwarg_name=kwarg_name,
231 )
234def returns_single_value(source_text: str, func: FunctionInfo) -> bool:
235 """Whether every ``return`` in the function returns a bare single value.
237 ``return`` / ``return x`` / ``return f(x)`` are single; ``return a, b`` /
238 ``return (a, b)`` / ``return [a, b]`` / ``return *a,`` are not. The
239 compiled boundary plan indexes into a returned sequence by
240 ``arg_index``, so a static single-tensor operator call can only replace
241 the output redistribution when the return is a single value.
243 The function's ``def`` span (decorators skipped) is re-parsed on its own;
244 a parse failure fails closed (``False``) — the caller demotes the static
245 form rather than guessing. The span starts at the ``def`` line's first
246 byte, which for a class method carries the method indent; ``dedent``
247 strips that common prefix so the isolated def parses as a top-level
248 statement (an indented first line would raise ``IndentationError`` and
249 fail every method closed).
250 """
251 if func.body_start is None or func.body_end is None:
252 return False
253 def_line_start = _def_keyword_offset(source_text, func)
254 span = textwrap.dedent(source_text[def_line_start:func.body_end])
255 try:
256 tree = ast.parse(span)
257 except SyntaxError:
258 return False
259 for node in ast.walk(tree):
260 if isinstance(node, ast.Return):
261 value = node.value
262 if value is None:
263 continue
264 if isinstance(value, (ast.Tuple, ast.List, ast.Starred)):
265 return False
266 return True
269def returns_any_value(source_text: str, func: FunctionInfo) -> bool:
270 """Whether every ``return`` exits with a single value or a tuple value.
272 Static tuple-output templates can apply output redistribution to
273 ``outputs[0]`` and rebuild the original object, matching the compiled plan's
274 index-0 behavior. A bare ``return`` has no value to transform; a list return
275 would need list-specific rebuilding, so it fails closed for now.
276 """
277 if func.body_start is None or func.body_end is None:
278 return False
279 def_line_start = _def_keyword_offset(source_text, func)
280 span = textwrap.dedent(source_text[def_line_start:func.body_end])
281 try:
282 tree = ast.parse(span)
283 except SyntaxError:
284 return False
285 seen_return = False
286 for node in ast.walk(tree):
287 if isinstance(node, ast.Return):
288 seen_return = True
289 if node.value is None:
290 return False
291 if isinstance(node.value, (ast.List, ast.Starred)):
292 return False
293 return seen_return
296def _def_keyword_offset(source_text: str, func: FunctionInfo) -> int:
297 """Byte offset of the ``def`` keyword line, skipping decorator lines.
299 ``func.def_offset`` anchors the FIRST line of the method (a decorator when
300 one is present); the span fed to ``ast.parse`` must start at the ``def``
301 keyword itself. Shared logic with ``_build_forward_impl_edit`` in
302 emit/parallel (kept local — three lines, not worth a cross-module helper).
303 """
304 def_line_start = func.def_offset
305 for line in source_text[def_line_start:func.body_start].splitlines(keepends=True):
306 if line.lstrip().startswith("def "):
307 break
308 def_line_start += len(line)
309 return def_line_start
312def _line_offsets(text: str) -> list[int]:
313 """Byte offset of each line's start (``offsets[line_no - 1]`` is 1-based line ``line_no``).
315 Single implementation, shared with ``codegen.source.compat`` (which imports
316 it from here — the existing layering direction is ``source`` -> ``astkit``).
317 """
318 offsets: list[int] = [0]
319 for line in text.splitlines(keepends=True):
320 offsets.append(offsets[-1] + len(line))
321 return offsets